"""Logic for line following.

UPDATE:(Vijay - 9/27/15) - Follower is outdated since we won't be
line-following for IEEE Hardware Competition 2016. This file has
been left here for reference;
Please do not instantiate it until it's fixed.
"""

from time import time
from time import sleep
import numpy as np

import bot.lib.lib as lib
import bot.hardware.ir_hub as ir_hub_mod
import bot.driver.mec_driver as mec_driver_mod
import pid as pid_mod
import bot.hardware.color_sensor as color_sensor_mod

from error_cases import LineLostError

class Follower(object):

    """Follows a line, detects intersections and stop conditions."""

    # Class variables
    # Array_Conditions
    Large_Object = 17  # the array sees a large object
    No_Line = 16  # the array see no line
    Noise = 19  # white values not next to each other

    # Variables for read_binary calls
    White_Black = True  # False  # True= white line, False= black line
    set_speed = 35

    THRESH = 60

    def __init__(self):
        # Build logger
        self.logger = lib.get_logger()

        # Build subsystems
        self.ir_hub = ir_hub_mod.IRHub()
        self.ir_hub.thrsh = 100
        self.driver = mec_driver_mod.MecDriver()
        self.color_sensor = color_sensor_mod.ColorSensor()

        # Build PIDs
        # FIXME 1 no longer in use
        self.front_right = pid_mod.PID()
        self.front_right_error = 0.0
        self.front_left = pid_mod.PID()
        self.front_left_error = 0.0
        self.back_right = pid_mod.PID()
        self.back_right_error = 0.0
        self.back_left = pid_mod.PID()
        self.back_left_error = 0.0
        self.strafe = pid_mod.PID()
        self.strafe_error = 0.0
        self.rotate_pid = pid_mod.PID()
        self.rotate_error = 0.0
        self.error = "NONE"

        # motor variables
        # FIXME 2 same as before
        self.translate_speed = 75
        self.prev_rate = 0

        # state variables
        self.heading = None  # must be initialize by caller
        self.front_state = Follower.No_Line
        self.back_state = Follower.No_Line
        self.left_state = Follower.No_Line
        self.right_state = Follower.No_Line

        # post error vars
        self.intersection = False
        self.lost_line = False
        self.timeLastUpdated = -1.0
        self.on_x = False

    @lib.api_call
    def get_translate_speed(self):
        return self.translate_speed

    @lib.api_call
    def set_translate_speed(self, speed):
        self.translate_speed = speed

    @lib.api_call
    def update(self):
        """Read IR values, compute aggregates."""
        ir_readings = self.ir_hub.read_binary(Follower.White_Black)
        for name, reading in ir_readings.iteritems():
            reading_arr = np.int_(reading)  # convert readings to numpy array
            np_reading = 4
            reading_sum = np.sum(np_reading)  # = no. of units lit
            if reading_sum > 0:
                self.ir_agg[name] = (
                    np.sum(self.ir_pos * reading_arr) / reading_sum)
            else:
                self.ir_agg[name] = None
        self.timeLastUpdated = time.time()

    @lib.api_call
    def get_strafe_error(self):
        return self.strafe_error

    @lib.api_call
    def get_rotate_error(self):
        return self.rotate_error

    def reset_errors(self):
        self.error = "NONE"
        # state variables
        # self.front_state = Follower.No_Line
        # self.back_state = Follower.No_Line
        # self.left_state = Follower.No_Line
        # self.right_state = Follower.No_Line

        # post error vars
        self.intersection = False
        self.lost_line = False
        self.timeLastUpdated = -1.0
        self.strafe_error = 0.0
        self.rotate_error = 0.0

        # reset pids
        self.strafe.clear_error()
        self.rotate_pid.clear_error()
        return

    @lib.api_call
    def report_states(self):
        # for debug of IR sensor state
        current_ir_reading = self.ir_hub.read_binary(Follower.White_Black)
        self.front_state = self.get_position_lr(
            current_ir_reading["front"])
        # Back is on the back side
        self.back_state = self.get_position_rl(
            current_ir_reading["back"])
        # Left is on the left
        self.left_state = self.get_position_lr(
            current_ir_reading["left"])
        # right is on the right
        self.right_state = self.get_position_rl(
            current_ir_reading["right"])
        self.logger.info("front = {}".format(self.front_state))
        self.logger.info("back = {}".format(self.back_state))
        self.logger.info("left = {}".format(self.left_state))
        self.logger.info("right = {}".format(self.right_state))

    @lib.api_call
    def oscillate(self, heading=0, osc_time=0.5):
        """Oscillate sideways, increasing in amplitude until line is found
        :param heading: The forward direction of the bot.
        :param osc_time: The initial time spent in each direction.

        """

        # Time in seconds for which bot oscillates in each direction.
        # Speed at which the bot is oscillating.
        # Increase in speed after each oscillation cycle.
        # Todo(Ahmed): Find reasonable constants.
        osc_speed = 10

        # The oscillation directions, perpendicular to parameter "heading"
        angle1 = heading + 90
        angle2 = heading - 90
        self.logger.debug(
            "Pre-correction angles: angle1: {}, angle2: {}".format(
                angle1, angle2))

        # Correct angles to fit bounds.
        angle1 %= self.driver.max_angle
        angle2 %= self.driver.max_angle
        self.logger.debug(
            "Post-correction angles: angle1: {}, angle2: {}".format(
                angle1, angle2))

        # Heading may be unecessary.
        # Test headings for valid 0,360 values.
        # assert 0 <= angle1 <= 360, "angle1 is {}".format(angle1)
        # assert 0 <= angle2 <= 360, "angle2 is {}".format(angle2)

        # Todo: Consider making this a function call.
        line_not_found = True
        while line_not_found:

            # Drives in each direction.
            self.driver.move(osc_speed, angle1)
            # Passes control to find line, which moves
            # until it finds line or runs out of time.
            # Note: watch_for_line returns "line_found"
            # (bool) and "time_elapsed" (int)
            results = self.watch_for_line(osc_time)
            self.driver.move(0, 0)

            if results["line_found"]:
                line_not_found = False

            # Search in other direction.
            self.driver.move(osc_speed, angle2)

            # "time elapsed" is used as max_time for more precise movements.
            results = self.watch_for_line(results["time_elapsed"] * 2)
            self.logger.debug(
                "Oscillation direction 1: osc_speed: {}, heading: {}".format(
                    osc_speed, heading))
            self.driver.move(0, 0)

            if results["line_found"]:
                line_not_found = False

            # If line is not found, Continue looping until line is found.
            # For now, stop when max speed is hit.
            osc_speed += 90
            if osc_speed >= self.driver.max_speed:
                line_not_found = False

    def reading_contains_pattern(self, pattern, reading):
        """Search the given reading for the given pattern.

        :param pattern: Pattern to search reading for.
        For example, [1, 1] for a pair of consecutive ones.
        :type pattern: list
        :param reading: IR array reading to search for the
        given pattern. Should contain only 0s and 1s.
        :type reading: list
        :returns: True if the pattern is in the reading, False otherwise.

        """
        return "".join(map(str, pattern)) in "".join(map(str, reading))

    def watch_for_line(self, max_time):
        """Recieves time period for which to continuously watch for line.
        Returns True when line is found.
        Returns False if line is not found before time is hit.
        """
        start_time = time()
        while True:
            reading = self.ir_hub.read_all()
            for name, array in reading.iteritems():
                if self.reading_contains_pattern([1, 1], array):
                    return {"line_found": True,
                            "time_elapsed": time() - start_time}
                if time() - start_time > max_time:
                    return {"line_found": False,
                            "time_elapsed": time() - start_time}

    @lib.api_call
    def assign_states(self, current_ir_reading=None):
        """ on_x=True flag does not allow intersection errors
            once left&right arrays clear intersection, on_x = false.
        Take 4x16 bit arrays and assigns the array to proper orientations.
        Note that the proper orientations are front, back, left and right.
        """
        # Keep prev front to ignore noise conditions
        # Keep prev back state to ignore large objects on back array
        prev_front_state = self.front_state
        prev_back_state = self.back_state
        # Get the current IR readings
        if current_ir_reading is None:
            current_ir_reading = self.ir_hub.read_binary(Follower.White_Black)

        # using heading, make front/back/left/right state assignments
        self.determine_states(current_ir_reading)

        # Clear on_x flag if off line on side arrays
        if(self.on_x and ((self.right_state > 15) or (self.left_state > 15))):
            self.on_x = False

        # Check for error conditions
        if((self.front_state > 15) or (self.back_state > 15) or
            (self.right_state < Follower.No_Line) and
                (self.left_state < Follower.No_Line)):

            # Lost Lines Superscede other conditions
            if((self.front_state == Follower.No_Line) and
                    (self.back_state == Follower.No_Line)):
                # Front and back lost line
                self.error = "LOST_LINE"
            elif(self.front_state == Follower.No_Line):
                # Front lost line
                self.error = "FRONT_LOST"
            elif(self.back_state == Follower.No_Line):
                # Back lost line
                self.error = "BACK_LOST"
            # Intersection preceds Large Object
            elif ((not (self.right_state == Follower.No_Line) and
                    not (self.left_state == Follower.No_Line)) and
                    not self.on_x):
                # Found Intersection because left and right lit up
                # if on_x=True, ignore this error
                self.error = "ON_INTERSECTION"
            elif((self.front_state == Follower.Large_Object)):
                # Found large object on front array.
                self.error = "LARGE_OBJECT"
            else:
                if(self.back_state == Follower.Large_Object):
                    # Ignore large objects on back array
                    # by using prev back state
                    self.back_state = prev_back_state
                if(self.front_state == Follower.Noise):
                    # Ignore Noise conditions
                    self.front_state = prev_front_state
                if(self.back_state == Follower.Noise):
                    # Ignore Noise conditions
                    self.back_state = prev_back_state

                    # self.error = "NONE"
        else:  # no errors
            self.error = "NONE"

        self.logger.info("FS: {}, BS {}, LS {}, RS {}".format(
            self.front_state,
            self.back_state,
            self.left_state,
            self.right_state))
        return self.front_state, self.back_state, \
            self.left_state, self.right_state

    def determine_states(self, current_ir_reading):
        if self.heading is None:
            self.heading = 180  # use implicit default value for testing
            self.logger.info("Using Test Heading = 180")
        # Heading east
        if self.heading == 270:
            # Forward is on the left side
            self.front_state = self.get_position_lr(
                current_ir_reading["left"])
            # Back is on the right side
            self.back_state = self.get_position_rl(
                current_ir_reading["right"])
            # Left is on the back
            self.left_state = self.get_position_lr(
                current_ir_reading["back"])
            # Right is on the front
            self.right_state = self.get_position_rl(
                current_ir_reading["front"])
        # Heading west
        elif self.heading == 90:
            # Forward is on the right side
            self.front_state = self.get_position_lr(
                current_ir_reading["right"])
            # Back is on the left side
            self.back_state = self.get_position_rl(
                current_ir_reading["left"])
            # Left is on the front
            self.left_state = self.get_position_lr(
                current_ir_reading["front"])
            # Right is on the back
            self.right_state = self.get_position_rl(
                current_ir_reading["back"])
        # Heading south
        elif self.heading == 180:
            # Forward is on the front side
            self.front_state = self.get_position_lr(
                current_ir_reading["front"])
            # Back is on the back side
            self.back_state = self.get_position_rl(
                current_ir_reading["back"])
            # Left is on the left
            # self.left_state = self.get_position_lr(
            #    current_ir_reading["left"])
            # right is on the right
            self.right_state = self.get_position_rl(
                current_ir_reading["right"])
            # for tri_state, copy right state to left
            self.left_state = self.right_state
        # Heading north
        elif self.heading == 0:
            # Forward is on the right side
            self.front_state = self.get_position_lr(
                current_ir_reading["back"])
            # Back is on the left side
            self.back_state = self.get_position_rl(
                current_ir_reading["front"])
            # Left is on the front
            self.left_state = self.get_position_lr(
                current_ir_reading["right"])
            # Right is on the back
            # self.right_state = self.get_position_rl(
            #    current_ir_reading["left"])
            # in tri state, copy left to right
            self.right_state = self.left_state

    def update_exit_state(self):
        if(self.error == "ON_INTERSECTION"):
            self.intersection = True
        elif(self.error == "LOST_LINE"):
            self.lost_line = True
        elif(self.error == "FRONT_LOST"):
            self.lost_line = True
        elif(self.error == "BACK_LOST"):
            self.lost_line = True

    def get_position_lr(self, readings):
        """Reading the IR sensors from left to right.

        Calculates the current state in reference to center from
        left to right. States go form -15 to 15.

        """
        self.hit_position = []
        state = 0.0
        for index, value in enumerate(readings):
            if(value == 1):
                self.hit_position.append(index)
        if len(self.hit_position) >= 4:
            # Error: Large Object detected
            return Follower.Large_Object
        if len(self.hit_position) == 0:
            # Error: No line detected
            return Follower.No_Line

        state = self.hit_position[0] * 2
        # Use first two hit irs to determine position on array
        # Ignores extra bits as noise.
        if len(self.hit_position) > 1:
            if self.hit_position[1] > 0:
                state = state + 1
            if abs(self.hit_position[0] - self.hit_position[1]) > 1:
                # Error: Discontinuity in sensors
                return Follower.Noise
        state = state - 15
        return state

    def get_position_rl(self, readings):
        """Reading the IR sensors from right to left.

        Calculates the current state in reference to center from
        right to left. States go form -15 to 15.

        """
        self.hit_position = []
        state = 0.0
        for index, value in enumerate(readings):
            if(value == 1):
                self.hit_position.append(index)
        if len(self.hit_position) >= 4:
            # Error: Large Object detected
            return Follower.Large_Object
        if len(self.hit_position) == 0:
            # Error: No line detected
            return Follower.No_Line

        state = self.hit_position[0] * 2
        # Use first two hit irs to determine position on array
        # Ignores extra bits as noise.
        if len(self.hit_position) > 1:
            if(self.hit_position[1] > 0):
                state = state + 1
            if(abs(self.hit_position[0] - self.hit_position[1]) > 1):
                # Error: Discontinuity in sensors
                return Follower.Noise
        state = (state - 15) * -1
        return state

    @lib.api_call
    def motors(self, bot_angle):
        """Used to update the motors speed and angular motion."""
        # If outside standard deviation angle, use rotate to straiten
        # if within std angle, use move to translate forward and strafe to
        # correct towards line horizontally
        # std is 7 on a range of -15 to 15
        std_angle = 7
        if abs(bot_angle) < std_angle:
            translate_angle = (self.strafe_error - self.heading + 180) % 360
            self.driver.move(self.translate_speed, translate_angle)
        else:
            # cap speed between (-100, 100)
            rotate_speed = max(-100, min(100, self.rotate_error))
            self.logger.info("rotate_speed = {}".format(rotate_speed))
            self.rotate(rotate_speed)

    @lib.api_call
    def smart_jerk(self):
        """Used to get the bot out of the box"""
        last_count = 0
        while True:
            count = 0
            ir_reading = self.ir_hub.read_binary(Follower.White_Black)
            for value in ir_reading["back"]:
                if(value == 1):
                    count += 1
            if((count != 0) and (last_count != 0)):
                self.driver.move(0, 0)
                return "DONE"
            self.driver.move(self.translate_speed, 0)
            self.logger.info("count = {}".format(count))
            last_count = count

    # Prevent rapid sign changes in rotate calls
    def rotate(self, rate):
        if(self.prev_rate * rate < 0):
            self.driver.move(0, 0)
            sleep(0.2)
        self.prev_rate = rate
        self.driver.rotate(rate)
        return

    @lib.api_call
    def center_on_intersection(self, heading=180):
        """center on intersection"""
        # first use center on line
        self.heading = heading
        self.center_on_line(heading)

        # then correct forwards/backwards
        forw_to_back_strafe = pid_mod.PID()
        # Init front_PID
        forw_to_back_strafe.set_k_values(5.5, 0.5, 2)
        while True:
            # kill momentum before reading
            self.driver.move(0, 0)
            # Assig states
            self.assign_states()
            # Check for error conditions
            if ((self.error != "NONE" and self.error != "ON_INTERSECTION")
                or (self.left_state == Follower.No_Line
                    or self.right_state == Follower.No_Line)):
                self.update_exit_state()
                self.logger.info("Error: {}".format(self.error))
                self.logger.info("FS: {}, BS: {}, lS: {}, RS: {}".format(
                    self.front_state,
                    self.back_state,
                    self.left_state,
                    self.right_state))
                self.driver.move(0, 0)
                return self.error
            # setup pid
            bot_front_position = (self.left_state + self.right_state) / 2
            # Call PID
            self.logger.info(
                "bot_front_position = {}".format(bot_front_position))
            position_error = forw_to_back_strafe.pid(
                0, bot_front_position, self.sampling_time)
            if(abs(bot_front_position) < 3):
                self.driver.move(0, 0)
                break
            # Cap at 0 and 100
            translate_speed = max(0, min(100, abs(position_error)))
            # use sign and heading to determin which direction to strafe
            if(position_error >= 0):
                translate_angle = (0 + self.heading) % 360
            else:
                translate_angle = (180 + self.heading) % 360
            if(abs(bot_front_position) < 5):
                return
            self.logger.info("translate_speed = {}".format(translate_speed))
            self.logger.info("translate_angle = {}".format(translate_angle))
            self.driver.move(translate_speed, translate_angle)
            # Take the current time set it equal to the previous time
        return "DONE"

    @lib.api_call
    def center_on_line(self, heading=180):
        """Used to center on the line"""
        center_rotate_pid = pid_mod.PID()
        side_to_side_strafe = pid_mod.PID()
        self.heading = heading
        while True:
            center_rotate_pid.clear_error()
            previous_time = time()
            # Init front_PID
            center_rotate_pid.set_k_values(6.75, .4, 1.75)
            while True:
                # kill momentum before reading
                self.driver.move(0, 0)
                # Assig states
                self.assign_states()
                # Check for error conditions
                if(self.error != "NONE" and self.error != "ON_INTERSECTION"):
                    self.update_exit_state()
                    self.logger.info("Error: {}".format(self.error))
                    self.logger.info("FS: {}, BS: {}, lS: {}, RS: {}".format(
                        self.front_state,
                        self.back_state,
                        self.left_state,
                        self.right_state))
                    self.driver.move(0, 0)
                    return self.error

                current_time = time()
                self.sampling_time = current_time - previous_time
                bot_angle = (self.front_state - self.back_state)
                # Call Rotate PID
                self.logger.info("bot_angle = {}".format(bot_angle))
                rotate_error = center_rotate_pid.pid(
                    0, bot_angle, self.sampling_time)
                # Report errors from strafe and rotate pid's
                if(abs(bot_angle) < 4):
                    self.driver.move(0, 0)
                    break
                self.logger.info("rotate_error = {}".format(rotate_error))
                rotate_speed = max(-100, min(100, rotate_error))
                self.rotate(rotate_speed)
                # Take the current time set it equal to the previous time
                previous_time = current_time

            side_to_side_strafe.clear_error()
            previous_time = time()
            # Init front_PID
            side_to_side_strafe.set_k_values(4.75, 2, 5.75)
            while True:
                # kill momentum before reading
                self.driver.move(0, 0)
                # Assig states
                self.assign_states()
                # Check for error conditions
                if(self.error != "NONE" and self.error != "ON_INTERSECTION"):
                    self.update_exit_state()
                    self.logger.info("Error: {}".format(self.error))
                    self.logger.info("FS: {}, BS: {}, lS: {}, RS: {}".format(
                        self.front_state,
                        self.back_state,
                        self.left_state,
                        self.right_state))
                    self.driver.move(0, 0)
                    return self.error
                # if angle off, break to rotate again
                bot_angle = (self.front_state - self.back_state)
                if(bot_angle >= 4):
                    break
                # calculate PID terms`
                current_time = time()
                bot_front_position = (self.front_state + self.back_state) / 2
                # Call side_to_side PID
                self.logger.info(
                    "bot_front_position = {}".format(bot_front_position))
                self.sampling_time = current_time - previous_time
                position_error = side_to_side_strafe.pid(
                    0, bot_front_position, self.sampling_time)
                # Report errors from strafe and rotate pid's
                if(abs(bot_front_position) < 3):
                    self.driver.move(0, 0)
                    break
                # Cap at 0 and 100
                translate_speed = max(0, min(100, abs(position_error)))
                # use sign and heading to determine which side to strafe to
                if(position_error >= 0):
                    translate_angle = (-90 + self.heading) % 360
                else:
                    translate_angle = (-270 + self.heading) % 360
                if(abs(bot_front_position) < 3):
                    return
                    self.logger.info(
                        "position_error = {}".format(
                            position_error))
                    self.logger.info(
                        "translate_speed = {}".format(
                            translate_speed))
                    self.logger.info(
                        "translate_angle = {}".format(
                            translate_angle))
                    self.driver.move(translate_speed, translate_angle)
                # Take the current time set it equal to the previous time
                previous_time = current_time

            if((abs(self.front_state) < 4) and (abs(self.back_state < 4))):
                return "Done"
        # end top while loop

    @lib.api_call
    def get_result(self):
        self.assign_states()
        return self.error

    @lib.api_call
    def strafe_to_line(self, heading=90):
        """Attempt to strafe sideways from one firing line to the next  """
        # assumes centered on line
        self.heading = heading
        self.assign_states()
        # strafe until off current line
        while ((self.left_state < Follower.No_Line) and
                (self.right_state < Follower.No_Line)):
            self.driver.move(self.translate_speed, heading)
            self.assign_states()
            self.driver.move(0, 0)
        # strafe until on next line
        while ((self.left_state == Follower.No_Line) or
                (self.right_state == Follower.No_Line)):
            self.driver.move(self.translate_speed, heading)
            self.assign_states()
            self.driver.move(0, 0)
        return "DONE STRAFING TO LINE"

    @lib.api_call
    def analog_state(self):
        """Make call to analog arrays"""

        # Take current time before reading ADC readings of the IRs
        previous_time = time()
        self.front_right.set_k_values(kp=.2, kd=0.1, ki=0.0)
        self.front_left.set_k_values(kp=.2, kd=0.1, ki=0.0)

        self.back_right.set_k_values(kp=.03, kd=0.009, ki=0.0)
        self.back_left.set_k_values(kp=.03, kd=0.009, ki=0.0)
        self.front_bin = [[0 for x in range(8)] for y in range(3)]

        self.front_right_error = 0.0
        self.front_left_error = 0.0
        self.back_right_error = 0.0
        self.back_left_error = 0.0

        while True:
            # Read ir arrays
            self.array_block = self.ir_hub.read_all()
            self.normalize_arrays()
            self.track_position()

            # Get the current time of the CPU
            current_time = time()
            self.sampling_time = current_time - previous_time

            # time before reading the IRs
            previous_time = time()

            # Count the number of hits
            front_hits = self.count_num_of_hits(self.array_block["front"])
            back_hits = self.count_num_of_hits(self.array_block["back"])

            self.front_bin[2] = self.front_bin[1]
            self.front_bin[1] = self.front_bin[0]
            self.front_bin[0] = self.assign_bin(self.array_block["front"])

            # print "fornt hits {}".format(self.front_bin[0])
            # right_hits = self.count_num_of_hits(self.array_block["right"])
            # left_hits = self.count_num_of_hits(self.array_block["left"])
            # print self.array_block

            # if not front_hits > 1:
            #    if((self.front_bin[1][0] == 1 \
            #        and self.front_bin[1][1] == 1) \
            #        or (self.front_bin[2][0] == 1 \
            #            and self.front_bin[2][1] == 1)):
            #        self.driver.move(0,0)
            #        return "left turn"

            #    if((self.front_bin[0][7] == 1 \
            #        and self.front_bin[1][7] == 1) \
            #        or (self.front_bin[0][6] == 1 \
            #        and self.front_bin[1][6] == 1)):
            #        self.driver.move(0,0)
            #        return "right turn"

            if front_hits > 3 or self.is_centerred_on_line('left') \
                    or self.is_centerred_on_line('right'):
                sleep(0.01)
                self.driver.move(60, 180)
                sleep(0.01)
                self.driver.move(0, 0)
                return "block or t-intersection"

            # else:
            #    if((self.front_bin[1][0] == 1 \
            #        and self.front_bin[1][1] == 1) \
            #        or (self.front_bin[2][0] == 1 \
            #            and self.front_bin[2][1] == 1)):
            #        self.driver.move(0,0)
            #        return "left turn at intersection"
            #    if((self.front_bin[0][7] == 1 \
            #        and self.front_bin[1][7] == 1) \
            #        or (self.front_bin[0][6] == 1 \
            #            and self.front_bin[1][6] == 1)):
            #        self.driver.move(0,0)
            #        return "right turn at intersection"

            if back_hits == 0 and front_hits == 0:
                self.driver.move(speed=0, angle=0)
                print "lost line"
                raise LineLostError("reading: {}".format(
                    self.array_block()))

            if(front_hits > 0):
                # Call PID
                self.front_right_error = (50 + self.front_right.pid(
                    0, self.bot_front_position, self.sampling_time))

                # Call PID
                self.front_left_error = (50 - self.front_left.pid(
                    0, self.bot_front_position, self.sampling_time))

                if(self.front_right_error >= 80):
                    self.front_right_error = 80
                elif(self.front_right_error <= -80):
                    self.front_right_error = -80

                if(self.front_left_error >= 80):
                    self.front_left_error = 80
                elif(self.front_left_error <= -80):
                    self.front_left_error = -80
            else:
                self.front_right_error = (50 + self.front_right.pid(
                    0, self.bot_back_position, self.sampling_time))

                # Call PID
                self.front_left_error = (50 - self.front_left.pid(
                    0, self.bot_back_position, self.sampling_time))

                if(self.front_right_error >= 80):
                    self.front_right_error = 80
                elif(self.front_right_error <= -80):
                    self.front_right_error = -80

                if(self.front_left_error >= 80):
                    self.front_left_error = 80
                elif(self.front_left_error <= -80):
                    self.front_left_error = -80

            if(back_hits > 0):
                # Call PID
                self.back_right_error = (50 + self.back_right.pid(
                    0, self.bot_back_position, self.sampling_time))

                # Call PID
                self.back_left_error = (50 - self.back_left.pid(
                    0, self.bot_back_position, self.sampling_time))

                if(self.back_right_error >= 80):
                    self.back_right_error = 80
                elif(self.back_right_error <= -80):
                    self.back_right_error = -80

                if(self.back_left_error >= 80):
                    self.back_left_error = 80
                elif(self.back_left_error <= -80):
                    self.back_left_error = -80
            else:
                # Call PID
                self.back_right_error = self.set_speed + \
                    self.back_right.pid(
                        0, self.bot_front_position, self.sampling_time)

                # Call PID
                self.back_left_error = self.set_speed - \
                    self.back_left.pid(
                        0, self.bot_front_position, self.sampling_time)

                if(self.back_right_error >= 80):
                    self.back_right_error = 80
                elif(self.back_right_error <= -80):
                    self.back_right_error = -80

                if(self.back_left_error >= 80):
                    self.back_left_error = 80
                elif(self.back_left_error <= -80):
                    self.back_left_error = -80

            self.driver.set_motor(name="front_right",
                                  value=self.front_right_error)
            self.driver.set_motor(name="front_left",
                                  value=self.front_left_error)
            self.driver.set_motor(name="back_right",
                                  value=self.back_right_error)
            self.driver.set_motor(name="back_left",
                                  value=self.back_left_error)

    def normalize_arrays(self):
        """Uesd to mormaliza ir readings coming form the ir array"""
        for array in self.array_block:
            for position, value in enumerate(self.array_block[array]):
                self.array_block[array][position] = (255 - value) - 100
                if(self.array_block[array][position] < 0):
                    self.array_block[array][position] = 0

    def track_position(self):
        """Trak the positon of the line"""
        self.bot_front_position = 0
        self.bot_back_position = 0
        value = max(self.array_block["front"])
        if not (value < 10):
            index = self.array_block["front"].index(value)
            # if value > 10:
            if index < 4:
                self.bot_front_position = 5.0 * (4.0 - index) \
                    * (4.0 - index) * (4.0 - index)
            else:
                self.bot_front_position = 5.0 * (3.0 - index) \
                    * (3.0 - index) * (3.0 - index)

        value = max(self.array_block["back"])
        if not(value < 10):
            index = self.array_block["back"].index(value)
            # print value
            # if value > 10:
            if index < 4:
                self.bot_back_position = 5.0 * (4.0 - index) \
                    * (4.0 - index) * (4.0 - index)
            else:
                self.bot_back_position = 5.0 * (3.0 - index) \
                    * (3.0 - index) * (3.0 - index)

    def count_num_of_hits(self, array):
        count = 0
        for value in array:
            if value > 90:
                count = count + 1
        return count

    def assign_bin(self, a_array):
        array = [0, 0, 0, 0, 0, 0, 0, 0]
        for index, value in enumerate(a_array):
            if value > self.THRESH:
                array[index] = 1
        return array

    @lib.api_call
    def check_for_branch(self, side):
        """Checks to see if there is a branch on the left.
        :returns: True or False
        """
        array_block = self.ir_hub.read_all()
        s_array = array_block[side]

        hits = self.count_num_of_hits(s_array)

        if hits <= 6:
            return True
        return False

    @lib.api_call
    def is_intersection(self):
        """Determines whether an anomaly is a turn or an intersection.
        """
        if (self.check_for_branch("right")
            or self.check_for_branch("left"))\
                and self.check_for_branch("front"):
            return True
        return False

    @lib.api_call
    def find_dir_of_int(self):
        if not (
                (self.check_for_branch("right")
                    or self.check_for_branch("left"))
                and self.check_for_branch("front")):
            return "error: not intersection"

        elif self.check_for_branch("right")\
                and self.check_for_branch("left"):
            return "error: branches in both dirs.111"

        elif self.check_for_branch("right"):
            return "right"

        elif self.check_for_branch("left"):
            return "left"

    @lib.api_call
    def find_dir_of_turn(self):
        """Determines whether a turn is on the right or left
        :returns: right, left, intersection, error
        """

        if (self.check_for_branch("right")
                and self.check_for_branch("left")):
            return "error: too many intersections"

        elif not self.check_for_branch("right") \
                and not self.check_for_branch("left"):
            return "error: no branches"

        elif self.check_for_branch("right"):
            return "right"

        elif self.check_for_branch("left"):
            return "left"

        else:
            return "error: No condition found. Line lost."

    @lib.api_call
    def is_centerred_on_line(self, side='front'):
        """Checks to see if bot is reasonably within center line.
        """
        arr_block = self.ir_hub.read_all()
        if arr_block[side][4] < self.THRESH \
                or arr_block[side][5] < self.THRESH:
            return True
        return False

    @lib.api_call
    def follow_ignoring_turns(self):
        while True:
            # self.recover()
            try:
                # Inch forward to be square on int
                # self.driver.jerk(speed=60,angle=0,duration=0.1)
                turn_dir = self.find_dir_of_turn()
                if turn_dir == "right" or turn_dir == "left":
                    self.rotate_to_line(turn_dir)
                    # self.driver.rough_rotate_90(turn_dir, r_time=0.6)
                    # Move out of intersection
                    # self.driver.drive(60, angle=0, duration=0.1)
                else:
                    self.recover()
                    break
            except LineLostError:
                self.logger.error("Line lost, attempting to recover")
                self.recover()

    @lib.api_call
    def rotate_to_line(self, direction, speed=50):
        """Rotates in given direction until line is found."""

        # Correct speed to match direction
        if direction == "right":
            speed = -speed

        self.driver.rotate(speed)
        sleep(0.3)  # allow time to leave current intersection

        while not self.is_centerred_on_line() \
                or self.check_for_branch('back'):
            self.logger.debug("looking for line")

        # throw in hard reverse to stop immediately
        self.driver.rotate(-speed)
        sleep(0.09)
        self.driver.rotate(0)

    @lib.api_call
    def drive_to_line(self, speed=50, angle=0):
        """Drives in certain direction blindly until line is found."""

        self.driver.move(speed, angle)
        while not self.is_centerred_on_line():
            self.logger.debug("looking for line")

        self.driver.hard_stop(speed, angle)

    @lib.api_call
    def recover(self):
        """When recovery is needed, this function continually
        calls itself until problem is solved.
        """
        # experiment with breaking any time front is detected.

        # if self.check_for_branch('front'):
        #    print "no need for recovery"
        #    return

        if not self.check_for_branch('front') \
                and not self.check_for_branch('back') \
                and not self.check_for_branch('left') \
                and not self.check_for_branch('right'):
            self.logger.debug("completely lost. Flailing")

        elif not self.check_for_branch('front') \
                and not self.check_for_branch('back') \
                and not self.check_for_branch('left') \
                and self.check_for_branch('right'):
            self.logger.debug("C1: right known: rotate right")
            self.rotate_to_line('right')
            self.recover()

        elif not self.check_for_branch('front') \
                and not self.check_for_branch('back') \
                and self.check_for_branch('left') \
                and not self.check_for_branch('right'):
            self.logger.debug("C2: left known: rotate left")
            self.rotate_to_line('left')
            self.recover()

        elif not self.check_for_branch('front') \
                and not self.check_for_branch('back') \
                and self.check_for_branch('left') \
                and self.check_for_branch('right'):
            self.logger.debug("C3: sides known: roate right")
            self.rotate_to_line('right')
            self.recover()

        elif not self.check_for_branch('front') \
                and self.check_for_branch('back') \
                and not self.check_for_branch('left') \
                and not self.check_for_branch('right'):
            self.logger.debug("C4: back known: Inch back")
            self.driver.drive(60, 180, 0.1)
            self.recover()

        elif not self.check_for_branch('front') \
                and self.check_for_branch('back') \
                and not self.check_for_branch('left') \
                and self.check_for_branch('right'):
            self.logger.debug("C5: back/right known: rotate right")
            self.rotate_to_line('right')

        elif not self.check_for_branch('front') \
                and self.check_for_branch('back') \
                and self.check_for_branch('left') \
                and not self.check_for_branch('right'):
            self.logger.debug("C6: back/left known: rotate right")
            self.rotate_to_line('left')

        # Todo: SHould this be a end condition?
        elif not self.check_for_branch('front') \
                and self.check_for_branch('back') \
                and self.check_for_branch('left') \
                and self.check_for_branch('right'):
            self.logger.debug("C7: front only known: inch fwd")
            self.driver.rough_rotate_90('right')
            self.driver.rough_rotate_90('right')
            self.recover()

        elif self.check_for_branch('front') \
                and not self.check_for_branch('back')  \
                and not self.check_for_branch('left')  \
                and not self.check_for_branch('right'):
            self.logger.debug("C8: Inch forward")
            self.driver.drive(60, 0, 0.1)
            self.recover()

        elif self.check_for_branch('front') \
                and not self.check_for_branch('back') \
                and not self.check_for_branch('left') \
                and self.check_for_branch('right'):
            # self.logger.debug("C9: front/side. inch towards")
            # self.driver.drive(60,90,0.1)
            # self.recover()
            return

        elif self.check_for_branch('front') \
                and not self.check_for_branch('back') \
                and self.check_for_branch('left') \
                and not self.check_for_branch('right'):
            # self.logger.debug("C10: shift left")
            # self.driver.drive(60,-90,0.1)
            # self.recover()
            return

        elif self.check_for_branch('front') \
                and self.check_for_branch('back') \
                and self.check_for_branch('left') \
                and self.check_for_branch('right'):
            self.logger.debug(
                "C15: Reached activity area. Take appropriate action.")
            return

        else:
            self.logger.debug(
                "Recovered successfully. Now following line - Cases 11 - 14")
            return


import os
import sys
from model import *
import numpy as np
import open3d as o3d
import copy
import rospy
import tensorflow as tf
import socket

from sensor_msgs.msg import PointCloud2, PointField
import sensor_msgs.point_cloud2 as pc2
import ctypes
import struct
from sensor_msgs.msg import PointCloud2, PointField
from std_msgs.msg import Header

BATCH_SIZE = 1
BATCH_SIZE_EVAL = 1
NUM_POINT = 4096
BASE_LEARNING_RATE = 0.001
GPU_INDEX = 0
MOMENTUM = 0.9
OPTIMIZER = 'adam'
DECAY_STEP = 300000
DECAY_RATE = 0.5

LOG_DIR = 'log'
if not os.path.exists(LOG_DIR): os.mkdir(LOG_DIR)
os.system('cp model.py %s' % (LOG_DIR)) # bkp of model def
os.system('cp train.py %s' % (LOG_DIR)) # bkp of train procedure
LOG_FOUT = open(os.path.join(LOG_DIR, 'log_train.txt'), 'w')

MAX_NUM_POINT = 4096
NUM_CLASSES = 2

BN_INIT_DECAY = 0.5
BN_DECAY_DECAY_RATE = 0.5
#BN_DECAY_DECAY_STEP = float(DECAY_STEP * 2)
BN_DECAY_DECAY_STEP = float(DECAY_STEP)
BN_DECAY_CLIP = 0.99

HOSTNAME = socket.gethostname()

def get_learning_rate(batch):
    learning_rate = tf.train.exponential_decay(
                        BASE_LEARNING_RATE,  # Base learning rate.
                        batch * BATCH_SIZE,  # Current index into the dataset.
                        DECAY_STEP,          # Decay step.
                        DECAY_RATE,          # Decay rate.
                        staircase=True)
    learning_rate = tf.maximum(learning_rate, 0.00001) # CLIP THE LEARNING RATE!!
    return learning_rate

def xyzrgb_array_to_pointcloud2(points, colors, stamp=None, frame_id=None, seq=None):
    '''
    Create a sensor_msgs.PointCloud2 from an array
    of points.
    '''

    points_ren = []
    lim = points.shape[0]
    print(points.shape)
    for k in range(lim):
        x = points[k,0]
        y = points[k,1]
        z = points[k,2]
        r = int(colors[k,0])
        g = int(colors[k,1])
        b = int(colors[k,2])
        a = int(255)

        rgb = struct.unpack('I', struct.pack('BBBB', b, g, r, a))[0]
        pt = [x, y, z, rgb]
        points_ren.append(pt)

    fields = [PointField('x', 0, PointField.FLOAT32, 1),
            PointField('y', 4, PointField.FLOAT32, 1),
            PointField('z', 8, PointField.FLOAT32, 1),
            PointField('rgba', 12, PointField.UINT32, 1),
            ]

    header = Header()
    header.frame_id = "camera_link"
    msg = pc2.create_cloud(header, fields, points_ren)
    return msg

class PointNet_Ros_Node:
  def __init__(self):
    '''initiliaze  ros stuff '''
    self.cloud_pub = rospy.Publisher("output/pointnet/segmented",PointCloud2)
    self.cloud_sub = rospy.Subscriber("pointnet_inference_cloud", PointCloud2,self.callback,queue_size=1)

    is_training = False
    batch = tf.Variable(0)
    learning_rate = get_learning_rate(batch)

    with tf.device('/gpu:'+str(GPU_INDEX)):
      pointclouds_pl, labels_pl = placeholder_inputs(BATCH_SIZE, NUM_POINT)
      is_training_pl = tf.placeholder(tf.bool, shape=())

      pred = get_model(pointclouds_pl, is_training_pl)
      loss = get_loss(pred, labels_pl)
      pred_softmax = tf.nn.softmax(pred)

      saver = tf.train.Saver()

    config = tf.ConfigProto(device_count = {'GPU': 0})
    config.gpu_options.allow_growth = True
    config.allow_soft_placement = True
    config.log_device_placement = True

    self.sess = tf.Session(config=config)
    tf.summary.scalar("loss", loss)
    merged = tf.summary.merge_all()

    optimizer = tf.train.AdamOptimizer(learning_rate)

    train_op = optimizer.minimize(loss, global_step=batch)

    self.ops = {'pointclouds_pl': pointclouds_pl,
    'labels_pl': labels_pl,
    'is_training_pl': is_training_pl,
    'pred': pred,
    'loss': loss,
    'train_op': train_op,
    'merged': merged,
    'step': batch}
    MODEL_PATH = "/home/atas/catkin_build_ws/src/ROS_NNs_FANUC_LRMATE200ID/dgcnn/tensorflow/sem_seg/log/model.ckpt"
    # Restore variables from disk.
    saver.restore(self.sess, MODEL_PATH)
    print("Model restored.")

  def callback(self, ros_point_cloud):
    xyz = np.array([[0,0,0]])
    rgb = np.array([[0,0,0]])
    #self.lock.acquire()
    gen = pc2.read_points(ros_point_cloud, skip_nans=True)
    int_data = list(gen)

    for x in int_data:
        test = x[3]
        # cast float32 to int so that bitwise operations are possible
        s = struct.pack('>f' ,test)
        i = struct.unpack('>l',s)[0]
        # you can get back the float value by the inverse operations
        pack = ctypes.c_uint32(i).value
        r = (pack & 0x00FF0000)>> 16
        g = (pack & 0x0000FF00)>> 8
        b = (pack & 0x000000FF)
        # prints r,g,b values in the 0-255 range
                    # x,y,z can be retrieved from the x[0],x[1],x[2]
        xyz = np.append(xyz,[[x[0],x[1],x[2]]], axis = 0)
        rgb = np.append(rgb,[[r,g,b]], axis = 0)
    self.eval_one_epoch(self.sess, self.ops,xyz,rgb)

  def eval_one_epoch(self,sess, ops, xyz, rgb):
    is_training = False
    test_writer = tf.summary.FileWriter(os.path.join(LOG_DIR, 'test'))

    xmax = 3.0
    xmin = -3.0
    current_data = np.zeros((4096,6))
    current_data[:,0:3]  = (xyz[0:NUM_POINT,:]- xmin) / (xmax  - xmin )
    current_data[:,3:6]  = rgb[0:NUM_POINT,:]/(255*255)

    current_data = current_data.reshape(1,4096, 6)
    current_label = np.zeros((1,4096))

    file_size = current_data.shape[0]
    num_batches = file_size // BATCH_SIZE_EVAL

    for batch_idx in range(num_batches):
      start_idx = batch_idx * BATCH_SIZE_EVAL
      end_idx = (batch_idx+1) * BATCH_SIZE_EVAL

      feed_dict = {ops['pointclouds_pl']: current_data[:, :],
                    ops['labels_pl']: current_label[:],
                    ops['is_training_pl']: is_training}
      summary, step, loss_val, pred_val = sess.run([ops['merged'], ops['step'], ops['loss'], ops['pred']],
                                    feed_dict=feed_dict)

      pred_label = np.argmax(pred_val, 2) # BxN

      test_writer.add_summary(summary, step)
      pred_val = np.argmax(pred_val, 2)
      correct = np.sum(pred_val == current_label[start_idx:end_idx])
      class_color = [[0,255,0],[0,0,255]]
      print(start_idx, end_idx)

      for i in range(start_idx, end_idx):
          print(pred_label.shape)
          pred = pred_label[i-start_idx, :]

          pts = current_data[i-start_idx, :, :]
          l = current_label[i-start_idx,:]

          xyz = np.zeros((NUM_POINT, 3))
          colors = np.zeros((NUM_POINT, 3))

          for j in range(NUM_POINT):
              l = int(current_label[i, j])

              pred_l = pred_val[i-start_idx, j]
              color = class_color[pred_l]
              color_gt = class_color[l]
              xyz[j, 0] ,a =  pts[j,0] *(xmax  - xmin ) + xmin,pts[j,0] *(xmax  - xmin ) + xmin
              xyz[j, 1] ,b =  pts[j,1] *(xmax  - xmin ) + xmin,pts[j,1] *(xmax  - xmin ) + xmin
              xyz[j, 2] ,c =  pts[j,2] *(xmax  - xmin ) + xmin,pts[j,2] *(xmax  - xmin ) + xmin
              xyz[j, 0] =  c
              xyz[j, 1] = -a
              xyz[j, 2] = -b

              colors[j,0],colors[j,1],colors[j,2] = color[0], color[1], color[2]

          #out_pcd = o3d.geometry.PointCloud()
          #out_pcd.points = o3d.utility.Vector3dVector(xyz)
          #out_pcd.colors = o3d.utility.Vector3dVector(colors)
          #o3d.io.write_point_cloud("/home/atas/predition.ply",out_pcd)
          msg = xyzrgb_array_to_pointcloud2(xyz,colors,rospy.Time.now(),"camera_link", 1)
          self.cloud_pub.publish(msg)

if __name__=='__main__':

  ic = PointNet_Ros_Node()
  rospy.init_node('ros_point_cloud',anonymous=True)
  rospy.spin()


import os
import glob
import json
import pickle
import yaml
import cv2
import numpy as np
from shapely.geometry import Polygon
from tqdm import tqdm

from extractor.common import Capture, delete_output
from extractor.gps import gps_to_ltp, gps_from_ltp
from extractor.keypoints import extract_keypoints, match_keypoints

def iuo(last_points, points):
    """Returns IoU of two polygons."""
    last_polygon = Polygon(last_points)
    polygon = Polygon(points)
    if not last_polygon.is_valid:
        return None
    if not polygon.is_valid:
        return None
    intersection = last_polygon.intersection(polygon).area
    union = polygon.union(last_polygon).area
    return intersection / union

def frame_overlap(last_pts, current_pts, image_size):
    """Returns the Intersection over Union of two frames based on matched keypoints."""
    img_h, img_w = image_size
    last_points = np.array([[0, 0],
                            [img_w, 0],
                            [img_w, img_h],
                            [0, img_h]], dtype=np.float32)
    transform, _ = cv2.findHomography(last_pts, current_pts, method=cv2.RANSAC)
    points = cv2.perspectiveTransform(last_points.reshape(-1, 1, 2), transform).reshape(-1, 2)
    iou = iuo(last_points, points)
    return iou

def make_camera_models_file(output_dir, camera_matrix, dist_coeffs, image_width, image_height):
    """Convert camera matrix and distortion coefficients into OpenSfM compatible format.
    Stores the camera_models_overrides file in the OpenSfM dataset."""
    k1, k2, p1, p2, k3 = dist_coeffs[0]

    scale = np.maximum(image_width, image_height)
    c_x = (camera_matrix[0, 2] - 0.5*image_width) / scale
    c_y = (camera_matrix[1, 2] - 0.5*image_height) / scale

    camera_models_overrides = {
        "all": {
            "projection_type": "brown",
            "width": image_width,
            "height": image_height,
            "focal_x": camera_matrix[0][0] / scale,
            "focal_y": camera_matrix[1][1] / scale,
            "c_x": c_x,
            "c_y": c_y,
            "k1": k1,
            "k2": k2,
            "k3": k3,
            "p1": p1,
            "p2": p2
        }
    }

    json.dump(camera_models_overrides, open(
        os.path.join(output_dir, "camera_models_overrides.json"), "w"))

def make_exif_file(output_dir, selected_frames, frame_names, gps, gps_dop, use_altitude):
    """Creates the exif_overrides file  for OpenSfM containing GPS positions of each selected frame."""
    exif_overrides = {}
    for selected_frame, frame_name in zip(selected_frames, frame_names):
        lon, lat, alt = gps[selected_frame]
        if not use_altitude:
            alt = 0.0
        exif_overrides["{}.jpg".format(frame_name)] = {
            "gps": {
                "latitude": lat,
                "longitude": lon,
                "altitude": alt,
                "dop": gps_dop
            }
        }
    json.dump(exif_overrides, open(os.path.join(output_dir, "exif_overrides.json"), "w"))

def select_frames_gps_visual(cap, gps, orb_detector, bf_matcher, frame_selection_gps_distance,
    frame_selection_visual_distance, match_distance_thres):
    """Select subset of frames based on travelled GPS distance and frame overlap (IoU).
    Returns indices of selected frames.

    Iterates over subsequent frames and computes the distance travelled along the GPS
    trajectory and a visual distance to the last keyframe. If either of the distances
    exceeds the specified threshold, the current frame is selected. The current frame
    is also used as the new keyframe for visual distance computation and the GPS
    distance is reset to zero. Visual distance is computed as 1.0 - IoU (intersection
    over union). IoU is estimated by finding the homography which maps the last frame
    onto the current frame and then computing the IoU between the frame boundaries of
    the last frame and the transformed current frame.
    """

    # for GPS distance computation
    gps_select = False
    last_pos = np.zeros((2,))
    gps_distance = 0

    # for visual distance computation
    visual_select = False
    last_kps = None
    last_des = None

    frame_idx = 0
    selected_frames = []
    pbar = tqdm(total=cap.num_images)
    while True:
        frame, _, _, _ = cap.get_next_frame(
            preprocess=True, undistort=True, equalize_hist=True)  # Note: undistort should probably be set to False
        if frame is None:
            break

        # test for GPS distance criterion
        pos = gps[frame_idx, :2]
        gps_distance += np.linalg.norm(pos - last_pos)
        last_pos = pos
        if gps_distance > frame_selection_gps_distance:
            gps_select = True

        # test for visual distance criterion
        kps, des = extract_keypoints(frame, orb_detector)
        visual_distance = np.inf
        if (last_kps is not None) and (len(kps) > 0):
            _, last_pts, current_pts = match_keypoints(
                bf_matcher, last_des, des, last_kps, kps,
                distance_threshold=match_distance_thres)

            if last_pts.shape[1] > 0 and current_pts.shape[1] > 0:

                # estimate motion with homography
                transform, _ = cv2.findHomography(
                    last_pts, current_pts, method=cv2.RANSAC)

                if transform is not None:
                    # compute frame overlap from transformed bounding rectangle
                    iou = frame_overlap(last_pts, current_pts, frame.shape)
                    if iou is not None:
                        visual_distance = 1.0 - iou

        if visual_distance > frame_selection_visual_distance:
            visual_select = True

        if gps_select or visual_select:
            gps_distance = 0
            last_kps = kps
            last_des = des
            gps_select = False
            visual_select = False
            selected_frames.append(frame_idx)

        frame_idx += 1
        pbar.update(1)
    pbar.close()

    return selected_frames

def select_frames_gps(gps, frame_selection_gps_distance):
    """Select subset of frames based on travelled GPS distance. Returns frame indices."""
    last_pos = np.zeros((2,))
    selected_frames = [0]
    travelled_distance = 0
    for i in range(len(gps)):
        pos = gps[i, :2]
        travelled_distance += np.linalg.norm(pos - last_pos)
        last_pos = pos
        if travelled_distance > frame_selection_gps_distance:
            travelled_distance = 0
            selected_frames.append(i)
    return selected_frames

def run(cluster, frames_root, calibration_root, output_dir, opensfm_settings,
        select_frames_mode, frame_selection_gps_distance, frame_selection_visual_distance,
        orb_nfeatures, orb_fast_thres, orb_scale_factor, orb_nlevels, match_distance_thres,
        gps_dop, output_video_fps):
    cluster_idx = cluster["cluster_idx"]
    frame_cluster = (cluster["frame_idx_start"], cluster["frame_idx_end"])

    output_dir = os.path.join(output_dir, "cluster_{:06d}".format(cluster_idx))
    delete_output(output_dir)
    os.makedirs(output_dir, exist_ok=True)

    camera_matrix = pickle.load(open(os.path.join(calibration_root, "camera_matrix.pkl"), "rb"))
    dist_coeffs = pickle.load(open(os.path.join(calibration_root, "dist_coeffs.pkl"), "rb"))

    # get video frames
    frame_files = sorted(glob.glob(os.path.join(frames_root, "radiometric", "*.tiff")))
    frame_files = frame_files[frame_cluster[0]:frame_cluster[1]]
    cap = Capture(frame_files, None, camera_matrix, dist_coeffs)

    make_camera_models_file(output_dir, camera_matrix, dist_coeffs, cap.img_w, cap.img_h)

    # write config.yaml for OpenSfM
    yaml.dump(opensfm_settings, open(os.path.join(output_dir, "config.yaml"), 'w'))

    # select frames based on travelled GPS trajectory
    gps = json.load(open(os.path.join(frames_root, "gps", "gps.json"), "r"))
    gps = np.array(gps[frame_cluster[0]:frame_cluster[1]])
    gps[:, 2] -= np.min(gps[:, 2])  # offset so minimum altitude is zero
    gps, origin = gps_to_ltp(gps)

    if select_frames_mode == "gps_visual":
        bf_matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
        orb_detector = cv2.ORB_create(
            nfeatures=orb_nfeatures, fastThreshold=orb_fast_thres,
            scaleFactor=orb_scale_factor, nlevels=orb_nlevels)
        selected_frames = select_frames_gps_visual(
            cap, gps, orb_detector, bf_matcher,
            frame_selection_gps_distance, frame_selection_visual_distance,
            match_distance_thres)
    elif select_frames_mode == "gps":
        selected_frames = select_frames_gps(gps, frame_selection_gps_distance)

    gps = gps_from_ltp(gps, origin)

    # store selected frames in OpenSfM dataset and write output video from selected frames for debugging
    video_path = os.path.join(output_dir, "selected_images.avi")
    fourcc = cv2.VideoWriter_fourcc(*"DIVX")
    videowriter = cv2.VideoWriter(video_path, fourcc, output_video_fps, (cap.img_w, cap.img_h))

    image_path = os.path.join(output_dir, "images")
    os.makedirs(image_path, exist_ok=True)

    frame_names = []
    for selected_frame in selected_frames:
        frame, _, frame_name, _ = cap.get_frame(selected_frame, preprocess=True, undistort=False, equalize_hist=True)  # Note: undistort should probably be set to False
        frame_names.append(frame_name)
        frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
        videowriter.write(frame)
        cv2.imwrite(os.path.join(output_dir, "images", "{}.jpg".format(frame_name)), frame)

    videowriter.release()

    # Create exif overwrites file
    make_exif_file(output_dir, selected_frames, frame_names, gps,
        gps_dop, opensfm_settings["use_altitude_tag"])


from chemutils.xyzutils import xyzToMolToRdkitMol
from chemutils.rdkitutils import getRdkitAtomXYZbyId, \
                                 rdkitMatchHeavyAtoms, \
                                 rdkitMolToInchi, \
                                 rdkitSubstrMatches, \
                                 rdkitSubstrAlignMatch, \
                                 setRdkitMolBondLength, \
                                 getRdkitMolBondLength, \
                                 rdkitMolToXYZ, \
                                 rdkitAlignMols
from chemutils.mathutils import findClosestXYZPoints, getXYZPointsDistance
from chemutils.ioutils import readFile, fileExists, removeBlankTrailingLines
from chemutils.obabelutils import obConvert
import random
import re
import numpy as np
import copy

def xyzToAtomsPositions(xyzFileOrStr):
    """
    Returns atom positions (order) given a molecule in an xyz format.
    Inchi-based algorithm.

    Note that the assigned positions are NOT UNIQUE and depend on the
    order of atoms in xyz file. This is true for both, heavy atoms
    and hydrogens. Heavy atoms positons are based on inchi, so they
    can only change due to presence of equivalent atoms with respect
    to connectivity. Positions of hydrogen atoms are assigned after
    the heavy atoms, taking any non-connected hydrogens first
    followed by hydrogens connected to heavy atoms.

    Use this function to set the atoms positions in a reference
    molecule. The idea is to assign the positions once and to never
    change them again.

    Arguments:
    ----------
    xyzFileOrStr : str
        input xyz molecule (either file path or xyz string)

    Returns:
    ----------
    atomsPositions: dict
        dictionary whose keys correspond to atoms positions in xyz
        file and values to the newly assigned positions
    """
    # get inchi with extra auxiliary log
    inchiWithAux = obConvert(inputMol=xyzFileOrStr,inputMolFormat='xyz',
                            outputMolFormat='inchi', options=['-xa'])
    inchi, inchiAux = inchiWithAux.split('\n')
    # find connectivity info in the inchi string - used to detect the
    # presence of heavy atoms.
    heavyAtomsInchiConnectivity = re.search(r'/c(\d+?\*)?(.*?/)',inchi)
    # read the mapping between heavy atoms (+ lone hydrogens) in xyz and inchi
    # from the auxiliary log
    heavyAtomsInchiAuxMap = re.search(r'/N:(.*?/)',inchiAux)
    # get the chemical formula from inchi - for fragments detection
    chemFormula = re.search(r'InChI=1S/(.*?/|\d*?H$)',inchi).groups()[0].replace('/','')
    # detect if the molecule contains any separate fragments
    molFragsDetect = re.search(r'(^\d+?[A-Z]|.+?\..+?)',chemFormula)

    # create the rdkit mol object which simply serves as the molecule container
    rdkitMolFromMol = xyzToMolToRdkitMol(xyzFileOrStr, removeHs=False)
    numAtoms = rdkitMolFromMol.GetNumAtoms()

    # initialise the atoms position dict
    atomsPositions = {k:None for k in range(numAtoms)}
    nextAtomId = 0

    # get the positions for heavy atoms
    if heavyAtomsInchiConnectivity is not None:
        # process the heavyAtomsInchiAuxMap and extract the atoms mapping
        heavyAtomsInchiAuxMap= heavyAtomsInchiAuxMap.groups()[0].replace('/','').replace(';',',').split(',')
        heavyAtomsMatch = {int(atomId)-1: i for i, atomId in enumerate(heavyAtomsInchiAuxMap)}

        if heavyAtomsMatch:
            # add the heavy atoms positions (+lone hydrogens) to the overall atomsPosition
            # dictionary
            atomsPositions = {**atomsPositions, **heavyAtomsMatch}
            nextAtomId = len(heavyAtomsMatch)
            hydrogensMap = {}
            # assign now positions of hydrogens that are attached to heavy atoms
            for heavyAtomId in heavyAtomsMatch.keys():
                heavyAtom = rdkitMolFromMol.GetAtomWithIdx(heavyAtomId)
                heavyAtomHsIds = [atom.GetIdx() for atom in heavyAtom.GetNeighbors() if atom.GetAtomicNum() == 1]
                # check the distance of each hydrogen to origin. I use it to select which hydrogen should have
                # the order assigned next. This way the hydrogens positions are less dependent on their xyz file order.
                atomDistFromOrigin = {}
                for atomId in heavyAtomHsIds:
                    atomPos = getRdkitAtomXYZbyId(rdkitMolFromMol,atomId)
                    atomDistFromOrigin[atomId] = getXYZPointsDistance(atomPos, np.zeros((3,)))
                # sort the hydrogens based on their distance to origin
                atomDistFromOriginSorted = {k: v for k, v in sorted(atomDistFromOrigin.items(), key=lambda item: item[1])}
                # add the hydrogen positions to the overall atoms positions
                for atomId in atomDistFromOriginSorted.keys():
                    hydrogensMap[atomId] = nextAtomId
                    nextAtomId += 1
                atomsPositions = {**atomsPositions, **hydrogensMap}

    # assign posititions to any atoms that are left
    # this step assigns order to molecules with no heavy atoms (e.g. h2 or h or h.h)
    if nextAtomId < numAtoms:
        nonHeavyAtomsIds = [atomId for atomId, refId in atomsPositions.items() if refId is None]
        nonHeavyAtomsMap = {}
        # similarly to hydrogens attached to heavy atoms, I order the remaining atoms with respect
        # to their distance to origin
        atomDistFromOrigin = {}
        for atomId in nonHeavyAtomsIds:
            atomPos = getRdkitAtomXYZbyId(rdkitMolFromMol,atomId)
            atomDistFromOrigin[atomId] = getXYZPointsDistance(atomPos, np.zeros((3,)))
        atomDistFromOriginSorted = {k: v for k, v in sorted(atomDistFromOrigin.items(), key=lambda item: item[1])}
        for atomId in atomDistFromOriginSorted.keys():
            nonHeavyAtomsMap[atomId] = nextAtomId
            nextAtomId += 1

        # add the remaining positions to the overall atoms positions
        atomsPositions = {**atomsPositions, **nonHeavyAtomsMap}

    # print a warning in case of xyz files containing separate fragments
    if molFragsDetect is not None:
        print('Warning: Provided xyz file contains two molecular fragments:',chemFormula)

    # check for duplicate and None values at the end
    hasDuplicates = len(atomsPositions.values()) > len(set(atomsPositions.values()))
    hasNones = None in atomsPositions.values()
    if hasDuplicates or hasNones:
        print('Error: atom canoncial positions algorithm has failed.')
        atomsPositions= {}
    return atomsPositions

def xyzMatch(xyzTargetFileOrStr, xyzRefFileOrStr):
    """
    Matches atom indices from one xyz molecule to the other.
    Molecules must have the same topology (the same inchis).
    Outputs match dictionary containing corresponding atoms
    indices between target (keys) and reference (values)
    molecules

    Arguments:
    ----------
    xyzTargetFileOrStr : str
        target molecule (either file path or xyz string)
    xyzRefFileOrStr : str
        reference molecule (either file path or xyz string)

    Returns:
    ----------
    match: dict
        dictionary whose keys correspond to target atoms
        and values to the matched reference atoms
    """
    # create ref and target molecules with hydrogens included
    rdkitRefMol = xyzToMolToRdkitMol(xyzRefFileOrStr, removeHs=False)
    rdkitTargetMol = xyzToMolToRdkitMol(xyzTargetFileOrStr, removeHs=False)

    # init the final atoms match dictionary to be populated
    # keys - target Mol atom ids, values - ref Mol atom ids
    # e.g. match = [
    #  0, 5  - target atom 3 matches ref atom 5
    #  1, 3  - target atom 1 matches ref atom 3
    #  ...
    # ]
    match = {k:None for k in range(rdkitTargetMol.GetNumAtoms())}

    # check if the two mols are the same
    if rdkitMolToInchi(rdkitRefMol) == rdkitMolToInchi(rdkitTargetMol):
        # find the best heavy atoms match first, e.g.
        # heavyAtomsMatch = {
        #  4: 2, - target heavy atom 4 matched ref heavy atom 2
        #  5: 6, - target heavy atom 5 matched ref heavy atom 6
        # }
        # I use dict here, because we only match a subset of atoms,
        # so we can not start from 0
        heavyAtomsMatch, _ = rdkitMatchHeavyAtoms(rdkitTargetMol, rdkitRefMol)

        # match any hydrogens attached to heavy atoms
        if heavyAtomsMatch:
            match = {**match, **heavyAtomsMatch}
            # now loop over heavy atoms in the target and ref mols,
            # get hydrogens bonded to them and match them based on their
            # distance with each other
            for targetAtomId, refAtomId in heavyAtomsMatch.items():
                targetAtom = rdkitTargetMol.GetAtomWithIdx(targetAtomId)
                targetAtomHsIds = [atom.GetIdx() for atom in targetAtom.GetNeighbors() if atom.GetAtomicNum() == 1]
                targetAtomHsPos = [getRdkitAtomXYZbyId(rdkitTargetMol,atomId) for atomId in targetAtomHsIds]

                refAtom = rdkitRefMol.GetAtomWithIdx(refAtomId)
                refAtomHsIds = [atom.GetIdx() for atom in refAtom.GetNeighbors() if atom.GetAtomicNum() == 1]
                refAtomHsPos = [getRdkitAtomXYZbyId(rdkitRefMol,atomId) for atomId in refAtomHsIds]

                closestHydrogens = findClosestXYZPoints(targetAtomHsPos, refAtomHsPos)

                # add matched hydrogen indices to the overall match dict
                for i, refId in enumerate(closestHydrogens):
                    # targetAtomHsIds[i] actual hydrogen id in the target molecule
                    # refAtomHsIds[refId] actual hydrogen id in the ref molecule
                    match[targetAtomHsIds[i]] = refAtomHsIds[refId]

        # check for any unmatched yet atoms
        # e.g any target atom ids whose ref ids are None
        unmatchedTargetAtomIds = [targetId for targetId, refId in match.items() if refId is None]
        if unmatchedTargetAtomIds:
            # get ref Ids not yet assigned to the target atoms
            unmatchedRefAtomIds = [refId for refId in range(rdkitTargetMol.GetNumAtoms()) if refId not in match.values()]

            unmatchedTargetAtomPos = [getRdkitAtomXYZbyId(rdkitTargetMol,atomId) for atomId in unmatchedTargetAtomIds]
            unmatchedRefAtomPos = [getRdkitAtomXYZbyId(rdkitRefMol,atomId) for atomId in unmatchedRefAtomIds]

            fragsMatch = findClosestXYZPoints(unmatchedTargetAtomPos, unmatchedRefAtomPos)
            # add matched hydrogen indices to the overall list
            for i, refId in enumerate(fragsMatch):
                match[unmatchedTargetAtomIds[i]] = unmatchedRefAtomIds[refId]
    else:
        print('Error: Provided xyz files contain two different molecules.')

    return match

def xyzReshuffle(xyzFileOrStr, seed):
    """
    Reshuffles the atoms order in xyz molecule string.

    Arguments:
    ----------
    xyzFileOrStr : str
        molecule (either file path or xyz string)

    Returns:
    ----------
    xyzStr: str
        xyz string containing reshuffled atoms
    """
    if fileExists(xyzFileOrStr): xyzFileOrStr= readFile(xyzFileOrStr)

    xyzFileOrStr = xyzFileOrStr.split('\n')
    # fix for xyz files containing extra empty lines at the end
    while not xyzFileOrStr[-1]: del xyzFileOrStr[-1]
    xyzHeader, xyzFileOrStr = xyzFileOrStr[:2], xyzFileOrStr[2:]
    random.Random(seed).shuffle(xyzFileOrStr)
    xyzFileOrStr = xyzHeader + xyzFileOrStr
    return "\n".join(xyzFileOrStr)

def xyzReorderOnAtomsMatch(xyzFileOrStr, atomsMatch):
    """
    Reorders the atoms xyz molecule string given
    the atoms match dictionary.

    Arguments:
    ----------
    xyzFileOrStr : str
        molecule (either file path or xyz string)

    Returns:
    ----------
    xyzStr: str
        xyz string containing re-ordered atoms
    """
    if fileExists(xyzFileOrStr): xyzFileOrStr= readFile(xyzFileOrStr)
    # take care of any trailing empty lines
    xyzFileOrStr = removeBlankTrailingLines(xyzFileOrStr)

    xyzFileOrStr = xyzFileOrStr.split('\n')
    xyzHeader, xyzFileOrStr = xyzFileOrStr[:2], xyzFileOrStr[2:]
    xyzReordered = [None]*len(xyzFileOrStr)
    for targetAtomId,refAtomId in atomsMatch.items():
        xyzReordered[refAtomId] = xyzFileOrStr[targetAtomId]
    xyzFileOrStr = xyzHeader + xyzReordered
    return "\n".join(xyzFileOrStr)

def xyzMatchWithBondAdjustment(xyzTargetFileOrStr, xyzRefFileOrStr, refAtomId1, refAtomId2, incrChange=2.5, nSteps=10):
    #
    # WORK IN PROGRESS, FEATURE NOT FINISHED
    #
    # create ref and target molecules with hydrogens included
    rdkitRefMol = xyzToMolToRdkitMol(xyzRefFileOrStr, removeHs=False)
    rdkitTargetMol = xyzToMolToRdkitMol(xyzTargetFileOrStr, removeHs=False)
    refInchi = rdkitMolToInchi(rdkitRefMol)
    tarInchi = rdkitMolToInchi(rdkitTargetMol)
    match = {}

    # get the chemical formula from inchi - for fragments detection
    refChemFormula = re.search(r'InChI=1S/(.*?/|\d*?H$)',refInchi).groups()[0].replace('/','')
    # detect if the molecule contains any separate fragments
    refMolFragsDetect = re.search(r'(^\d+?[A-Z]|.+?\..+?)',refChemFormula)

    rdkitRefMolLocal = copy.deepcopy(rdkitRefMol)
    rdkitTarMolLocal = copy.deepcopy(rdkitTargetMol)
    startBondLength = getRdkitMolBondLength(rdkitRefMolLocal, refAtomId1, refAtomId2)
    bestRmsd = 1e9
    bestBondLength = startBondLength
    for step in range(nSteps):
        refInchiLocal = rdkitMolToInchi(rdkitRefMolLocal)
        refChemFormulaLocal = re.search(r'InChI=1S/(.*?/|\d*?H$)',refInchiLocal).groups()[0].replace('/','')
        refMolFragsDetectLocal = re.search(r'(^\d+?[A-Z]|.+?\..+?)',refChemFormulaLocal)
        if not refMolFragsDetectLocal: break

        newBondLength = startBondLength-incrChange*(step+1)
        setRdkitMolBondLength(rdkitRefMolLocal, refAtomId1, refAtomId2, newBondLength)
        rmsd = rdkitAlignMols(rdkitRefMolLocal, rdkitTargetMol)
        if rmsd < bestRmsd:
            bestRmsd = rmsd
            bestBondLength = newBondLength

    for step in range(nSteps):
        newBondLength = startBondLength+incrChange*(step+1)
        setRdkitMolBondLength(rdkitRefMolLocal, refAtomId1, refAtomId2, newBondLength)

        refInchi2 = rdkitMolToInchi(rdkitRefMolLocal)
        tarInchi = rdkitMolToInchi(rdkitTargetMol)

        bb = rdkitRefMolLocal.GetNumAtoms()
        aa = rdkitTargetMol.GetNumAtoms()
        #rmsd = rdkitAlignMols(rdkitTarMolLocal, rdkitRefMolLocal)
        match = rdkitSubstrAlignMatch(rdkitTargetMol, rdkitRefMolLocal)
        match = rdkitMatchHeavyAtoms(rdkitTargetMol, rdkitRefMolLocal)
        if rmsd < bestRmsd:
            bestRmsd = rmsd
            bestBondLength = newBondLength

    match= xyzMatch(xyzTargetFileOrStr, xyzRefFileOrStr)


from cluster_env import Cluster
from RL_brain import QLearningTable
import datetime
import numpy as np
import matplotlib.pyplot as plt
import pylab as pl
import random
import pandas as pd
import time
from pandas.testing import assert_frame_equal

def state_init():
    init_state = pd.DataFrame(np.zeros(327*8).reshape(327, 8), columns=[0, 1, 2, 3, 4, 5, 6, 7])
    for i in range(len(init_state)):
        j = random.randint(0, 7)
        init_state.iloc[i][j] = 1
    return init_state

def server_attribute():
    init_state = pd.DataFrame(np.zeros(8*24).reshape(8, 24), columns=np.arange(24))
    server = []
    attribute_list = range(len(init_state.columns))
    for i in range(len(init_state)):
        attribute_list = list(set(attribute_list).difference(set(server)))
        server = random.sample(attribute_list, 3)
        for j in server:
            init_state.iloc[i][j] = 1
    return init_state

def update():
    cost_all_list = []
    reward_all_list = []
    query_number = len(env.QSs)
    # print(query_number)
    for episode in episodes:
        epoch_curr_time1 = datetime.datetime.now()
        # initial state
        state_init_arr = env.state_array(env.state_init)
        state = (env.state_init).copy()
        costs = env.cost_init
        sum = 0
        reward_list = [0]
        state_arr_for_one = state_init_arr
        reward = init_reward
        while True:
            # RL choose action based on observation
            # The action here is a number(the index of the real action)
            action = RL.choose_action(str(state))
            # print(action)

            # RL take action and get next observation and reward
            state_, costs_, reward_, cost_all, is_better = env.step(action, state, costs)

            state_arr = env.state_array(state_)
            different = [y for y in (state_arr_for_one + state_arr) if y not in state_arr_for_one]
            print("diffrent:", different)
            state_arr_for_one = state_arr
            different_init = [y for y in (state_init_arr + state_arr) if y not in state_init_arr]

            if ((reward_ < init_reward and reward_ < min(reward_list) or
                (len(different) == 0 and reward_ >= reward and reward_ > (init_reward)))):
                done = True
            else:
                done = False
            # RL learn from this transition
            print("done:", done)

            if ((reward_ < init_reward) and (reward_ < min(reward_list))):
                print("reward值小于初始值或并且该循环的最小值")

            if len(different) == 0 and reward_ >= reward and reward_ > (init_reward*1.3):
                print("reward值大于前一个循环的reward值并且采取动作后状态不改变")

            if reward_ < init_reward:
                reward_ = -1

            reward = reward_

            reward_list.append(reward)

            RL.learn(str(state), action, reward, str(state_), done)

            costs = costs_

            # q_table=RL.q_table.copy()
            state = state_

            sum += 1
            # print(sum)

            if done:
                break

        reward_all_list.append(reward)
        epoch_curr_time2 = datetime.datetime.now()
        epoch_time = epoch_curr_time2 - epoch_curr_time1
                # if (action in actions and q_table.loc[str(state), action] >= 0) and (done and q_table.loc[str(state), action] >= 0 and reward > 0):
                #     break
                # else:
                #     actions.append(action)
                # break while loop when end of this episode

                # if done and q_table.loc[str(state),action]!=0:
                #     break

        cost_all_list.append(cost_all)
        print("epoch:", episode+1)
        print("The number of cycles in this epoch：", sum)
        print("The reward list:", reward_list)
        print("The best reward in this epoch：", max(reward_list))
        print("The final reward in this epoch:", reward)
        print("The final cost in this epoch:", cost_all)
        print("当前状态与初始状态的差别", (different_init))
        print("当前状态与初始状态的差别数", len(different_init))
        print("epoch_time:", epoch_time, "\n")
    # end of game
    # print("final state:\n",state)
    print("------------------------")
    print("The final state_array:", state_arr)
    print("The final cost:", cost_all)
    # if state.equals(env.state_init):
    #     print("True")
    # else:
    #     print("False")
        # assert_frame_equal(state,env.state_init)
    # env.destroy()
    return cost_all_list, reward_all_list

if __name__ == "__main__":
    improve_list = []
    test_number = 50
    state_init = state_init()
    # for i in range(50):
    # print("第%d次测试：" % (i+1))
    episodes = np.arange(100)
    curr_time1 = datetime.datetime.now()

    # print(len(state_init))
    server_attribute = pd.DataFrame(np.array([0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                                              0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0,
                                              0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0,
                                              1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0,
                                              0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
                                              0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0,
                                              0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0,
                                              0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0]).
                                    reshape(8, 24),
                                    columns=np.arange(24))

    env = Cluster(state_init, server_attribute)
    init_reward = env.reward(env.cost_all(env.cost_init), env.state_init)
    RL = QLearningTable(actions=list(range(env.n_actions)))
    # env.after(100, update)
    # env.mainloop()
    cost_all_list, reward_all_list = update()
    curr_time2 = datetime.datetime.now()
    train_time = curr_time2-curr_time1
    print("The training time：", train_time)
    print("\n")
    improve = ((reward_all_list[-1] - init_reward)/init_reward)*100
    print("The improve percent:", improve, "%")
    improve_list.append(improve)

    y_1 = reward_all_list
    y_all_list = y_1
    x = episodes
    y = y_all_list
    y1 = [init_reward]*len(episodes)
    fig = plt.Figure(figsize=(7, 5))
    pl.plot(x, y, label=u'RL')
    pl.legend()
    pl.plot(x, y1, label=u'Init')
    pl.legend()
    pl.xlabel(u"epoch", size=14)
    pl.ylabel(u"reward", size=14)
    plt.show()
    # y = improve_list
    # x = np.arange(test_number)
    # fig = plt.Figure(figsize=(7, 5))
    # pl.plot(x, y, label=u'Improve')
    # pl.legend()
    # plt.show()


import os.path
import sys
import math
import argparse
import time
import random
from collections import OrderedDict,deque
import matplotlib
matplotlib.use('Agg')
# matplotlib.use('Qt5Agg')
import torch
import numpy as np
import options.options as option
from utils import util
from data import create_dataloader, create_dataset
from models import create_model
from utils.logger import Logger, PrintLogger
import tqdm
from datetime import datetime
import cv2
import copy

# USE_Y_GENERATOR_4_CHROMA = True

def main():
    # options
    parser = argparse.ArgumentParser()
    parser.add_argument('-opt', type=str, required=True, help='Path to option JSON file.')
    parser.add_argument('-single_GPU', action='store_true',help='Utilize only one GPU')
    parser.add_argument('-chroma', action='store_true',help='Training the chroma-channels generator')
    if parser.parse_args().single_GPU:
        available_GPUs = util.Assign_GPU(maxMemory=0.66)
    else:
        # available_GPUs = util.Assign_GPU(max_GPUs=None,maxMemory=0.8,maxLoad=0.8)
        available_GPUs = util.Assign_GPU(max_GPUs=None)
    opt = option.parse(parser.parse_args().opt, is_train=True,batch_size_multiplier=len(available_GPUs),name='JPEG'+('_chroma' if parser.parse_args().chroma else ''))

    if not opt['train']['resume']:
        util.mkdir_and_rename(opt['path']['experiments_root'])  # Modify experiment name if exists
        util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' and \
            not key == 'pretrained_model_G' and not key == 'pretrained_model_D'))
    option.save(opt)
    opt = option.dict_to_nonedict(opt)  # Convert to NoneDict, which return None for missing key.

    # print to file and std_out simultaneously
    sys.stdout = PrintLogger(opt['path']['log'])

    # random seed
    seed = opt['train']['manual_seed']
    if seed is None:
        seed = random.randint(1, 10000)
    print("Random Seed: ", seed)
    random.seed(seed)
    torch.manual_seed(seed)

    # create train and val dataloader
    for phase, dataset_opt in opt['datasets'].items():
        if phase == 'train':
            max_accumulation_steps = max([opt['train']['grad_accumulation_steps_G'], opt['train']['grad_accumulation_steps_D']])
            train_set = create_dataset(dataset_opt)
            train_size = int(math.ceil(len(train_set) / dataset_opt['batch_size']))
            print('Number of train images: {:,d}, iters: {:,d}'.format(len(train_set), train_size))
            total_iters = int(opt['train']['niter']*max_accumulation_steps)#-current_step
            total_epoches = int(math.ceil(total_iters / train_size))
            print('Total epoches needed: {:d} for iters {:,d}'.format(total_epoches, total_iters))
            train_loader = create_dataloader(train_set, dataset_opt)
        elif phase == 'val':
            val_dataset_opt = dataset_opt
            val_set = create_dataset(dataset_opt)
            val_loader = create_dataloader(val_set, dataset_opt)
            print('Number of val images in [{:s}]: {:d}'.format(dataset_opt['name'], len(val_set)))
        else:
            raise NotImplementedError('Phase [{:s}] is not recognized.'.format(phase))
    assert train_loader is not None
    DEBUG = False
    # Create model
    if DEBUG:
        from models.base_model import BaseModel
        model = BaseModel
        model.step = 0
    else:
        model = create_model(opt,max_accumulation_steps,chroma_mode=opt['name'][:len('JPEG/chroma')]=='JPEG/chroma')

    # create logger
    logger = Logger(opt)
    # Save validation set results as image collage:
    SAVE_IMAGE_COLLAGE = True
    start_time,start_time_gradient_step = time.time(),model.step // max_accumulation_steps
    save_GT_Uncomp = True
    lr_too_low = False
    print('---------- Start training -------------')
    last_saving_time = time.time()
    recently_saved_models = deque(maxlen=4)
    for epoch in range(int(math.floor(model.step / train_size)),total_epoches):
        for i, train_data in enumerate(train_loader):
            model.gradient_step_num = model.step // max_accumulation_steps
            not_within_batch = model.step % max_accumulation_steps == (max_accumulation_steps - 1)
            saving_step = ((time.time()-last_saving_time)>60*opt['logger']['save_checkpoint_freq']) and not_within_batch
            if saving_step:
                last_saving_time = time.time()

            # save models
            if lr_too_low or saving_step:
                model.save_log()
                recently_saved_models.append(model.save(model.gradient_step_num))
                if len(recently_saved_models)>3:
                    model_2_delete = recently_saved_models.popleft()
                    os.remove(model_2_delete)
                    if model.D_exists:
                        os.remove(model_2_delete.replace('_G.','_D.'))
                print('{}: Saving the model before iter {:d}.'.format(datetime.now().strftime('%H:%M:%S'),model.gradient_step_num))
                if lr_too_low:
                    break

            if model.step > total_iters:
                break

            # time_elapsed = time.time() - start_time
            # if not_within_batch:    start_time = time.time()
            # log
            if model.gradient_step_num % opt['logger']['print_freq'] == 0 and not_within_batch:
                logs = model.get_current_log()
                print_rlt = OrderedDict()
                print_rlt['model'] = opt['model']
                print_rlt['epoch'] = epoch
                print_rlt['iters'] = model.gradient_step_num
                # time_elapsed = time.time() - start_time
                print_rlt['time'] = (time.time() - start_time)/np.maximum(1,model.gradient_step_num-start_time_gradient_step)
                start_time, start_time_gradient_step = time.time(), model.gradient_step_num
                for k, v in logs.items():
                    print_rlt[k] = v
                print_rlt['lr'] = model.get_current_learning_rate()
                logger.print_format_results('train', print_rlt,keys_ignore_list=['avg_est_err'])
                model.display_log_figure()

            # validation
            if (not_within_batch or i==0) and (model.gradient_step_num) % opt['train']['val_freq'] == 0: # and model.gradient_step_num>=opt['train']['D_init_iters']:
                print_rlt = OrderedDict()
                if model.generator_changed:
                    print('---------- validation -------------')
                    start_time = time.time()
                    if False and SAVE_IMAGE_COLLAGE and model.gradient_step_num%opt['train']['val_save_freq'] == 0: #Saving training images:
                        # GT_image_collage,quantized_image_collage = [],[]
                        cur_train_results = model.get_current_visuals(entire_batch=True)
                        train_psnrs = [util.calculate_psnr(util.tensor2img(cur_train_results['Decomp'][im_num], out_type=np.uint8,min_max=[0,255]),
                            util.tensor2img(cur_train_results['Uncomp'][im_num], out_type=np.uint8,min_max=[0,255])) for im_num in range(len(cur_train_results['Decomp']))]
                        #Save latest training batch output:
                        save_img_path = os.path.join(os.path.join(opt['path']['val_images']),
                                                     '{:d}_Tr_PSNR{:.3f}.png'.format(model.gradient_step_num, np.mean(train_psnrs)))
                        util.save_img(np.clip(np.concatenate((np.concatenate([util.tensor2img(cur_train_results['Uncomp'][im_num], out_type=np.uint8,min_max=[0,255]) for im_num in
                                 range(len(cur_train_results['Decomp']))],0), np.concatenate(
                                [util.tensor2img(cur_train_results['Decomp'][im_num], out_type=np.uint8,min_max=[0,255]) for im_num in range(len(cur_train_results['Decomp']))],
                                0)), 1), 0, 255).astype(np.uint8), save_img_path)
                    Z_latent = [0]+([-0.5,0.5] if opt['network_G']['latent_input'] else [])
                    print_rlt['psnr'] = 0
                    for cur_Z in Z_latent:
                        model.perform_validation(data_loader=val_loader,cur_Z=cur_Z,print_rlt=print_rlt,GT_and_quantized=save_GT_Uncomp,
                                                 save_images=((model.gradient_step_num) % opt['train']['val_save_freq'] == 0) or save_GT_Uncomp)
                    if save_GT_Uncomp:  # Save GT Uncomp images
                        save_GT_Uncomp = False
                    print_rlt['psnr'] /= len(Z_latent)
                    model.log_dict['psnr_val'].append((model.gradient_step_num,print_rlt['psnr']))
                else:
                    print('Skipping validation because generator is unchanged')
                # time_elapsed = time.time() - start_time
                # Save to log
                print_rlt['model'] = opt['model']
                print_rlt['epoch'] = epoch
                print_rlt['iters'] = model.gradient_step_num
                # print_rlt['time'] = time_elapsed
                print_rlt['time'] = (time.time() - start_time)/np.maximum(1,model.gradient_step_num-start_time_gradient_step)
                # model.display_log_figure()
                # model.generator_changed = False
                logger.print_format_results('val', print_rlt,keys_ignore_list=['avg_est_err'])
                print('-----------------------------------')

            model.feed_data(train_data,mixed_Y=True)
            model.optimize_parameters()

            # update learning rate
            if not_within_batch:
                lr_too_low = model.update_learning_rate(model.gradient_step_num)
            # current_step += 1
        if lr_too_low:
            print('Stopping training because LR is too low')
            break

    print('Saving the final model.')
    model.save(model.gradient_step_num)
    model.save_log()
    print('End of training.')

if __name__ == '__main__':
    # # OpenCV get stuck in transform when used in DataLoader
    # # https://github.com/pytorch/pytorch/issues/1838
    # # However, cause problem reading lmdb
    # import torch.multiprocessing as mp
    # mp.set_start_method('spawn', force=True)
    main()