import datetime
from ev3sim.visual.objects import visualFactory
from ev3sim.settings import ObjectSetting
import numpy as np
import pymunk
import pygame
from ev3sim.search_locations import batch_locations, preset_locations
from ev3sim.simulation.interactor import IInteractor
from ev3sim.simulation.loader import ScriptLoader
from ev3sim.simulation.randomisation import Randomiser
from ev3sim.simulation.world import World, stop_on_pause
from ev3sim.objects.base import objectFactory, STATIC_CATEGORY
from ev3sim.objects.utils import local_space_to_world_space, magnitude_sq
from ev3sim.file_helper import find_abs
from ev3sim.visual.manager import ScreenObjectManager
from ev3sim.visual.utils import screenspace_to_worldspace

class RescueInteractor(IInteractor):

    # Must occur before device interactors.
    SORT_ORDER = -10

    FOLLOW_POINT_CATEGORY = 0b1000
    SHOW_FOLLOW_POINTS = True
    SHOW_ROBOT_COLLIDER = False
    FOLLOW_POINT_COLLISION_TYPE = 5
    ROBOT_CENTRE_COLLISION_TYPE = 6
    ROBOT_CENTRE_RADIUS = 3
    FOLLOW_POINT_RADIUS = 1
    # You can be at most this far from the previous follow point before lack of progress is called.
    MAX_FOLLOW_DIST = 8
    # Randomisation constants
    BOT_SPAWN_RADIUS = 2
    BOT_SPAWN_ANGLE = [-10, 10]

    BOT_SPAWN_POSITION = [[[0, 0], 2]]
    TILE_DEFINITIONS = [
        {
            "path": "tiles/definitions/city_limits.yaml",
            "position": [-2, -2],
            "rotation": 0,
            "flip": False,
        },
    ]

    GAME_LENGTH_MINUTES = 5

    TILE_LENGTH = 30
    _pressed = False
    _touches = 0
    _touch_points = 0

    CAN_SPAWN_POSITION = (0, 0)

    MODE_FOLLOW = "FOLLOW"
    MODE_RESCUE = "RESCUE"

    # Don't autostart in the `startUp` parent call.
    AUTOSTART_BOTS = False

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.time_tick = 0
        self.hover_rect = visualFactory(
            name="Rectangle",
            width=30,
            height=30,
            position=(0, 0),
            fill=None,
            stroke="#ff0000",
            stroke_width=1,
            zPos=20,
        )
        self.hover_rect.key = "hover_rect"

    def spawnTiles(self):
        self.tiles = []
        for i, tile in enumerate(self.TILE_DEFINITIONS):
            self.tiles.append({})
            import yaml

            path = find_abs(tile["path"], allowed_areas=preset_locations())
            with open(path, "r") as f:
                t = yaml.safe_load(f)
            self.maxZpos = 0
            base_pos = np.array(tile.get("position", [0, 0]))
            # Transfer to rescue space.
            base_pos = [base_pos[0] * self.TILE_LENGTH, base_pos[1] * self.TILE_LENGTH]
            base_rotation = tile.get("rotation", 0) * np.pi / 180
            flip = tile.get("flip", False)
            for obj in t["elements"]:
                rel_pos = np.array(obj.get("position", [0, 0]))
                if flip:
                    rel_pos[0] = -rel_pos[0]
                    if obj.get("name", "") == "Image":
                        obj["flip"] = [True, False]
                    if obj.get("name", "") == "Arc":
                        obj["rotation"] = 180 - obj.get("rotation", 0)
                        obj["angle"] = -obj["angle"]
                obj["rotation"] = (obj.get("rotation", 0)) * np.pi / 180 + base_rotation
                obj["position"] = local_space_to_world_space(rel_pos, base_rotation, base_pos)
                obj["sensorVisible"] = True
                k = obj["key"]
                obj["key"] = f"Tile-{i}-{k}"
                self.maxZpos = max(self.maxZpos, obj.get("zPos", 0))
            t["elements"].append(
                {
                    "position": local_space_to_world_space(np.array([0, 0]), base_rotation, base_pos),
                    "rotation": base_rotation,
                    "type": "visual",
                    "name": "Rectangle",
                    "width": self.TILE_LENGTH,
                    "height": self.TILE_LENGTH,
                    "fill": None,
                    "stroke_width": 0.1,
                    "stroke": "rescue_outline_color",
                    "zPos": self.maxZpos + 0.1,
                    "key": f"Tile-{i}-outline",
                    "sensorVisible": False,
                }
            )
            self.tiles[-1]["type"] = t.get("type", "follow")
            self.tiles[-1]["follows"] = []
            self.tiles[-1]["roam_status"] = []
            self.tiles[-1]["world_pos"] = base_pos
            self.tiles[-1]["rotation"] = base_rotation
            self.tiles[-1]["flip"] = flip
            if self.tiles[-1]["type"] == "follow":
                self.tiles[-1]["entries"] = t["entries"]
                self.tiles[-1]["exits"] = t["exits"]
                mname, cname = t.get("checker").rsplit(".", 1)
                import importlib

                klass = getattr(importlib.import_module(mname), cname)
                with open(find_abs(t["ui"], allowed_areas=preset_locations()), "r") as f:
                    self.tiles[-1]["ui_elem"] = yaml.safe_load(f)
                for j, point in enumerate(t["follow_points"]):
                    if isinstance(point[0], (list, tuple)):
                        self.tiles[-1]["follows"].append([])
                        self.tiles[-1]["roam_status"].append([])
                        for path in point:
                            self.tiles[-1]["follows"][-1].append([])
                            self.tiles[-1]["roam_status"][-1].append([])
                            for point2 in path:
                                if isinstance(point2, str):
                                    self.tiles[-1]["roam_status"][-1][-1][-1] = point2
                                else:
                                    if flip:
                                        point2[0] = -point2[0]
                                    self.tiles[-1]["follows"][-1][-1].append(
                                        local_space_to_world_space(
                                            np.array(point2), tile.get("rotation", 0) * np.pi / 180, base_pos
                                        )
                                    )
                                    self.tiles[-1]["roam_status"][-1][-1].append(None)
                    else:
                        if isinstance(point, str):
                            self.tiles[-1]["roam_status"][-1] = point
                        else:
                            if flip:
                                point[0] = -point[0]
                            self.tiles[-1]["follows"].append(
                                local_space_to_world_space(
                                    np.array(point), tile.get("rotation", 0) * np.pi / 180, base_pos
                                )
                            )
                            self.tiles[-1]["roam_status"].append(None)
                self.tiles[-1]["checker"] = klass(self.tiles[-1]["follows"], i, self, **t.get("checker_kwargs", {}))
            else:
                self.tiles[-1]["green_conns"] = t.get("green_conns", [])
            self.tiles[-1]["all_elems"] = ScriptLoader.instance.loadElements(t["elements"])
        connecting_objs = []
        for tile in self.tiles:
            # We need to add connections between green tiles if they exist.
            if tile["type"] == "follow":
                continue
            under, right, under_right = self._checkConnectingRescueTiles(tile["world_pos"])
            if under_right and under and right:
                # Draw a big square connecting all 4 tiles.
                key = "c1-" + str(tile["world_pos"][0]) + "-" + str(tile["world_pos"][1])
                connecting_objs.append(
                    {
                        "type": "visual",
                        "name": "Rectangle",
                        "width": 50,
                        "height": 50,
                        "fill": "grass_color",
                        "zPos": 0.3,
                        "key": key,
                        "position": [
                            tile["world_pos"][0] + self.TILE_LENGTH / 2,
                            tile["world_pos"][1] + self.TILE_LENGTH / 2,
                        ],
                        "sensorVisible": True,
                    }
                )
            else:
                if under:
                    key = "c2-" + str(tile["world_pos"][0]) + "-" + str(tile["world_pos"][1])
                    connecting_objs.append(
                        {
                            "type": "visual",
                            "name": "Rectangle",
                            "width": 20,
                            "height": 50,
                            "fill": "grass_color",
                            "zPos": 0.3,
                            "key": key,
                            "position": [tile["world_pos"][0], tile["world_pos"][1] + self.TILE_LENGTH / 2],
                            "sensorVisible": True,
                        }
                    )
                if right:
                    key = "c3-" + str(tile["world_pos"][0]) + "-" + str(tile["world_pos"][1])
                    connecting_objs.append(
                        {
                            "type": "visual",
                            "name": "Rectangle",
                            "width": 50,
                            "height": 20,
                            "fill": "grass_color",
                            "zPos": 0.3,
                            "key": key,
                            "position": [tile["world_pos"][0] + self.TILE_LENGTH / 2, tile["world_pos"][1]],
                            "sensorVisible": True,
                        }
                    )
        self.connecting_objs = ScriptLoader.instance.loadElements(connecting_objs)

    def _checkConnectingRescueTiles(self, world_pos, sx=1, sy=1):
        under = False
        right = False
        under_right = False
        for tile in self.tiles:
            if tile["type"] == "follow":
                continue
            if (
                tile["world_pos"][0] - world_pos[0] == self.TILE_LENGTH * sx
                and tile["world_pos"][1] - world_pos[1] == 0
            ):
                right = True
            if (
                tile["world_pos"][0] - world_pos[0] == 0
                and tile["world_pos"][1] - world_pos[1] == self.TILE_LENGTH * sy
            ):
                under = True
            if (
                tile["world_pos"][0] - world_pos[0] == self.TILE_LENGTH * sx
                and tile["world_pos"][1] - world_pos[1] == self.TILE_LENGTH * sy
            ):
                under_right = True
        return under, right, under_right

    def _recurseObj(self, obj, indicies):
        for index in indicies:
            obj = obj[index]
        return obj

    def lackOfProgress(self, bot_index):
        print("Lack of Progress!")
        World.instance.paused = True
        self.current_follow = None

    def collidedFollowPoint(self, follow_indexes, bot_index):
        if self._recurseObj(self.follow_completed[follow_indexes[0]], follow_indexes[1:]):
            return
        roam = self._recurseObj(self.tiles[follow_indexes[0]]["roam_status"], follow_indexes[1:])
        if roam is not None:
            if roam == "ROAM_START":
                self.roaming[bot_index] = (True, follow_indexes[1:])
                self.roam_tile[bot_index] = follow_indexes[0]
            if roam == "ROAM_END":
                self.roaming[bot_index] = False
        else:
            if (
                self.roaming[bot_index]
                and follow_indexes[1:] > self.roaming[bot_index][1]
                and follow_indexes[0] == self.roam_tile[bot_index]
            ):
                # We should hit the roam end trigger first.
                self.lackOfProgress(bot_index)
        self.current_follow = follow_indexes
        self._recurseObj(self.follow_completed[follow_indexes[0]], follow_indexes[1:-1])[follow_indexes[-1]] = True
        self._recurseObj(self.tiles[follow_indexes[0]]["follow_colliders"], follow_indexes[1:]).visual.fill = "#00ff00"
        self.tiles[follow_indexes[0]]["checker"].onNewFollowPoint(self.follow_completed[follow_indexes[0]])

    def follow_point_colour(self, indicies):
        return (
            "#ff0000" if self._recurseObj(self.tiles[indicies[0]]["roam_status"], indicies[1:]) is None else "#0000ff"
        )

    def _spawnFollowAtLocationWithIndicies(self, position, indicies):
        key = "Tile-follow-" + "-".join(list(map(str, indicies)))
        obj = objectFactory(
            **{
                "collider": "inherit",
                "visual": {
                    "name": "Circle",
                    "radius": self.FOLLOW_POINT_RADIUS,
                    "fill": self.follow_point_colour(indicies) if self.SHOW_FOLLOW_POINTS else None,
                    "stroke_width": 0,
                    "sensorVisible": False,
                    "zPos": self.maxZpos + 0.2,
                },
                "position": position,
                "physics": True,
                "static": True,
                "key": key,
            }
        )
        obj.shape.filter = pymunk.ShapeFilter(categories=self.FOLLOW_POINT_CATEGORY)
        obj.shape.sensor = True
        obj.shape._follow_indexes = indicies
        obj.shape.collision_type = self.FOLLOW_POINT_COLLISION_TYPE
        World.instance.registerObject(obj)
        if self.SHOW_FOLLOW_POINTS:
            ScreenObjectManager.instance.registerObject(obj, obj.key)
        return obj

    def spawnFollowPointPhysics(self):
        for i, tile in enumerate(self.tiles):
            tile["follow_colliders"] = []
            for j, pos in enumerate(tile["follows"]):
                if isinstance(pos[0], (list, tuple)):
                    tile["follow_colliders"].append([])
                    for k, path in enumerate(pos):
                        tile["follow_colliders"][-1].append([])
                        for l, pos2 in enumerate(path):
                            obj = self._spawnFollowAtLocationWithIndicies(pos2, (i, j, k, l))
                            tile["follow_colliders"][-1][-1].append(obj)
                else:
                    obj = self._spawnFollowAtLocationWithIndicies(pos, (i, j))
                    tile["follow_colliders"].append(obj)

    TILE_UI_ELEM_HEIGHT = 10
    TILE_UI_PADDING = 20
    TILE_UI_INITIAL_HEIGHT = 98

    @property
    def tileUIHeight(self):
        return self.TILE_UI_PADDING * 2 + self.TILE_UI_ELEM_HEIGHT * len(
            [tile for tile in self.tiles if tile["type"] == "follow"]
        )

    def spawnTileUI(self):
        elems = []
        i = 0
        for tile in self.tiles:
            if tile["type"] == "follow":
                elem = tile["ui_elem"]
                elem["key"] = f"Tile-{i}-UI"
                elem["position"] = [
                    -140,
                    -self.TILE_UI_ELEM_HEIGHT
                    * (i - (len([tile for tile in self.tiles if tile["type"] == "follow"]) - 1) / 2),
                ]
                elems.append(elem)
                i += 1
        for i, spawned in enumerate(ScriptLoader.instance.loadElements(elems)):
            self.tiles[i]["ui_spawned"] = spawned
            self.tiles[i]["checker"].onSpawn()
        self.can_complete_box = visualFactory(
            name="Rectangle",
            width=6,
            height=6,
            fill="UI_tilebox_not_complete",
            stroke="#ffffff",
            stroke_width=0.1,
            zPos=5,
            position=(-138, -self.tileUIHeight / 2 + self.TILE_UI_PADDING / 2 + 3),
        )
        self.line_reentry_box = visualFactory(
            name="Rectangle",
            width=6,
            height=6,
            fill="UI_tilebox_not_complete",
            stroke="#ffffff",
            stroke_width=0.1,
            zPos=5,
            position=(-128, -self.tileUIHeight / 2 + self.TILE_UI_PADDING / 2 + 3),
        )
        ScreenObjectManager.instance.registerVisual(self.can_complete_box, "can_complete")
        ScreenObjectManager.instance.registerVisual(self.line_reentry_box, "line_reentry")
        ScriptLoader.instance.object_map["rescueBGMid"].scale = (1, self.tileUIHeight / self.TILE_UI_INITIAL_HEIGHT)
        ScriptLoader.instance.object_map["rescueBGMid"].calculatePoints()
        ScriptLoader.instance.object_map["rescueBGTop"].position = (-146.6, self.tileUIHeight / 2)
        ScriptLoader.instance.object_map["rescueBGBottom"].position = (-146.6, -self.tileUIHeight / 2)
        ScriptLoader.instance.object_map["rescueScoreSum"].position = (
            -110,
            -self.tileUIHeight / 2 + self.TILE_UI_PADDING / 2,
        )
        ScriptLoader.instance.object_map["touchesTitle"].position = (
            ScriptLoader.instance.object_map["touchesTitle"].position[0],
            self.tileUIHeight / 2 - self.TILE_UI_PADDING / 2,
        )
        ScriptLoader.instance.object_map["touchesCount"].position = (
            ScriptLoader.instance.object_map["touchesCount"].position[0],
            self.tileUIHeight / 2 - self.TILE_UI_PADDING / 2,
        )
        ScriptLoader.instance.object_map["touchesScore"].position = (
            ScriptLoader.instance.object_map["touchesScore"].position[0],
            self.tileUIHeight / 2 - self.TILE_UI_PADDING / 2,
        )
        self.touchesChanged()

    def spawnCan(self):
        can_vis = {
            "type": "object",
            "physics": True,
            "visual": {
                "name": "Circle",
                "radius": 2.5,
                "fill": "#e85d04",
                "stroke_width": 0.1,
                "stroke": "#eae2b7",
                "zPos": 5,
            },
            "position": self.CAN_SPAWN_POSITION,
            "key": "can",
            "friction": 0.7,
            "restitution": 0.2,
        }
        self.can_obj = ScriptLoader.instance.loadElements([can_vis])[0]

    def spawnRobotFollows(self):
        self.bot_follows = []
        for i in range(len(self.robots)):
            # Spawn the robot follow point collider.
            obj = objectFactory(
                **{
                    "collider": "inherit",
                    "visual": {
                        "name": "Circle",
                        "radius": self.ROBOT_CENTRE_RADIUS,
                        "fill": "#00ff00" if self.SHOW_ROBOT_COLLIDER else None,
                        "stroke_width": 0,
                        "sensorVisible": False,
                        "zPos": 100,
                    },
                    "physics": True,
                    "key": f"Robot-{id}-follow",
                }
            )
            obj.shape.filter = pymunk.ShapeFilter(categories=self.FOLLOW_POINT_CATEGORY)
            obj.shape.sensor = True
            obj.shape.collision_type = self.ROBOT_CENTRE_COLLISION_TYPE
            obj.shape._robot_index = i
            World.instance.registerObject(obj)
            if self.SHOW_ROBOT_COLLIDER:
                ScreenObjectManager.instance.registerObject(obj, obj.key)
            self.bot_follows.append(obj)

    def startUp(self):
        super().startUp()
        self.START_TIME = datetime.timedelta(minutes=self.GAME_LENGTH_MINUTES)
        self.spawnTiles()
        self.spawnFollowPointPhysics()
        self.spawnTileUI()
        self.spawnRobotFollows()
        assert len(self.robots) <= len(self.BOT_SPAWN_POSITION), "Not enough spawning locations specified."
        self.spawnCan()

        self.reset()
        for i in range(len(self.robots)):
            # This is bad, I should get the robot key somehow else (Generally speaking the robot class, interactor and object should be more tightly coupled.)
            self.bot_follows[i].body.position = [
                float(v)
                for v in local_space_to_world_space(
                    ScriptLoader.instance.robots[f"Robot-{i}"]._follow_collider_offset,
                    self.robots[i].body.angle,
                    (self.robots[i].body.position.x, self.robots[i].body.position.y),
                )
            ]
        self.addCollisionHandler()

        for robot in self.robots:
            robot.robot_class.onSpawn()

    def addCollisionHandler(self):
        handler = World.instance.space.add_collision_handler(
            self.FOLLOW_POINT_COLLISION_TYPE, self.ROBOT_CENTRE_COLLISION_TYPE
        )

        saved_world_no = World.instance.spawn_no

        def handle_collide(arbiter, space, data):
            if World.instance.spawn_no != saved_world_no:
                return
            a, b = arbiter.shapes
            if hasattr(a, "_follow_indexes"):
                self.collidedFollowPoint(a._follow_indexes, b._robot_index)
            elif hasattr(b, "_follow_indexes"):
                self.collidedFollowPoint(b._follow_indexes, a._robot_index)
            else:
                raise ValueError("Two objects with collision types used by rescue don't have a tile follow point.")
            return False

        handler.begin = handle_collide

    def spawnAt(self, tileIndex):
        self.can_scored_this_spawn = False
        self.resetFollows()
        spawn_point = (
            self.tiles[tileIndex]["follows"][0][0][0]
            if isinstance(self.tiles[tileIndex]["follows"][0], (list, tuple))
            else self.tiles[tileIndex]["follows"][0]
        )
        for i in range(len(self.robots)):
            self.robots[i].body.angle = self.tiles[tileIndex]["rotation"]
            if self.tiles[tileIndex]["flip"]:
                self.robots[i].body.angle = self.robots[i].body.angle + np.pi
            self.bot_follows[i].body.position = [
                float(v)
                for v in local_space_to_world_space(
                    ScriptLoader.instance.robots[f"Robot-{i}"]._follow_collider_offset,
                    self.robots[i].body.angle,
                    (self.robots[i].body.position.x, self.robots[i].body.position.y),
                )
            ]
            self.robots[i].body.position = [
                a + b - c
                for a, b, c in zip(self.robots[i].body.position, spawn_point, self.bot_follows[i].body.position)
            ]
            self.robots[i].body.velocity = (0, 0)
            self.robots[i].body.angular_velocity = 0
        self.touchBot()

    def resetFollows(self):
        self.current_follow = None
        self.follow_completed = [
            [[[False for point in path] for path in x] if isinstance(x, (list, tuple)) else False for x in y["follows"]]
            for y in self.tiles
        ]
        for x in range(len(self.tiles)):
            for y in range(len(self.tiles[x]["follow_colliders"])):
                if isinstance(self.tiles[x]["follow_colliders"][y], (list, tuple)):
                    for z, path in enumerate(self.tiles[x]["follow_colliders"][y]):
                        for w in range(len(path)):
                            path[w].visual.fill = self.follow_point_colour((x, y, z, w))
                else:
                    self.tiles[x]["follow_colliders"][y].visual.fill = self.follow_point_colour((x, y))
        self.roaming = [False] * len(self.robots)
        self.roam_tile = [-1] * len(self.robots)

    def reset(self):
        self.modes = [self.MODE_FOLLOW] * len(self.robots)
        self.resetPositions()
        self.time_tick = 0
        self._touches = 0
        self._touch_points = 0
        self.touchesChanged()
        self.setScore(0)
        self.resetFollows()
        self.resetCanUI()
        for x in range(len(self.tiles)):
            if self.tiles[x]["type"] == "follow":
                self.tiles[x]["checker"].onReset()
        for rob_id in ScriptLoader.instance.robots:
            ScriptLoader.instance.startProcess(rob_id)

    def setScore(self, val):
        self.score = val
        ScriptLoader.instance.object_map["rescueScoreSum"].text = str(self.score)

    def incrementScore(self, val):
        self.setScore(self.score + val)

    def decrementScore(self, val):
        self.setScore(self.score - val)

    def resetCanUI(self):
        self.can_scored_this_spawn = False
        self.can_scored = False
        self.can_obj.body.position = self.CAN_SPAWN_POSITION
        self.can_obj.position = np.array(self.can_obj.body.position)
        self.reentered = False
        self.canStateChanged()

    def canStateChanged(self):
        self.can_complete_box.fill = "#00ff00" if self.can_scored else "#000000"
        self.line_reentry_box.fill = "#00ff00" if self.reentered else "#000000"

    def scoreCan(self):
        self.can_scored_this_spawn = True
        if not self.can_scored:
            self.can_scored = True
            self.incrementScore(50)
            self.canStateChanged()

    def scoreReEntry(self):
        if not self.reentered:
            self.reentered = True
            self.incrementScore(20)
            self.canStateChanged()

    def resetPositions(self):
        for i, robot in enumerate(self.robots):
            diff_radius = Randomiser.random() * self.BOT_SPAWN_RADIUS
            diff_angle = Randomiser.random() * 2 * np.pi
            robot.body.position = [
                self.BOT_SPAWN_POSITION[i][0][0] * self.TILE_LENGTH + diff_radius * np.cos(diff_angle),
                self.BOT_SPAWN_POSITION[i][0][1] * self.TILE_LENGTH + diff_radius * np.sin(diff_angle),
            ]
            robot.body.angle = (
                (
                    self.BOT_SPAWN_POSITION[i][1]
                    + self.BOT_SPAWN_ANGLE[0]
                    + Randomiser.random() * (self.BOT_SPAWN_ANGLE[1] - self.BOT_SPAWN_ANGLE[0])
                )
                * np.pi
                / 180
            )
            robot.body.velocity = (0.0, 0.0)
            robot.body.angular_velocity = 0

    def tick(self, tick):
        super().tick(tick)
        self.cur_tick = tick
        for i in range(len(self.robots)):
            self.bot_follows[i].body.position = [
                float(v)
                for v in local_space_to_world_space(
                    ScriptLoader.instance.robots[f"Robot-{i}"]._follow_collider_offset,
                    self.robots[i].body.angle,
                    (self.robots[i].body.position.x, self.robots[i].body.position.y),
                )
            ]
            # Ensure visual is not 1 frame behind.
            self.bot_follows[i].position = np.array(self.bot_follows[i].body.position)
            if self.current_follow is not None and not World.instance.paused:
                distance = magnitude_sq(
                    self.bot_follows[i].position
                    - self._recurseObj(self.tiles[self.current_follow[0]]["follows"], self.current_follow[1:])
                )
                if distance > self.MAX_FOLLOW_DIST * self.MAX_FOLLOW_DIST and not self.roaming[i]:
                    self.lackOfProgress(i)
            if self.roaming[i] and not World.instance.paused:
                # Check we are still in the tile. Kinda bad.
                if (
                    abs(self.bot_follows[i].position[0] - self.tiles[self.roam_tile[i]]["world_pos"][0])
                    > self.TILE_LENGTH / 2
                    or abs(self.bot_follows[i].position[1] - self.tiles[self.roam_tile[i]]["world_pos"][1])
                    > self.TILE_LENGTH / 2
                ) and magnitude_sq(
                    self.bot_follows[i].position
                    - self._recurseObj(self.tiles[self.roam_tile[i]]["follows"], self.roaming[i][1])
                ) > (
                    self.ROBOT_CENTRE_RADIUS + self.FOLLOW_POINT_RADIUS + 0.05
                ) * (
                    self.ROBOT_CENTRE_RADIUS + self.FOLLOW_POINT_RADIUS + 0.05
                ):
                    self.lackOfProgress(i)
            if not World.instance.paused:
                x = (self.bot_follows[i].position[0] + self.TILE_LENGTH / 2) // self.TILE_LENGTH
                y = (self.bot_follows[i].position[1] + self.TILE_LENGTH / 2) // self.TILE_LENGTH
                for tile in self.tiles:
                    tx = (tile["world_pos"][0] + self.TILE_LENGTH / 2) // self.TILE_LENGTH
                    ty = (tile["world_pos"][1] + self.TILE_LENGTH / 2) // self.TILE_LENGTH
                    if tx == x and ty == y:
                        if tile["type"] == "rescue" and self.modes[i] == self.MODE_FOLLOW:
                            self.modes[i] = self.MODE_RESCUE
                            # For reentry.
                            self.resetFollows()
                            self.current_follow = None
                        elif tile["type"] == "follow" and self.modes[i] == self.MODE_RESCUE:
                            self.modes[i] = self.MODE_FOLLOW
                            if self.can_scored_this_spawn:
                                self.scoreReEntry()
                        break
                else:
                    self.lackOfProgress(i)
                if not self.can_scored_this_spawn:
                    cx = (self.can_obj.position[0] + self.TILE_LENGTH / 2) // self.TILE_LENGTH
                    cy = (self.can_obj.position[1] + self.TILE_LENGTH / 2) // self.TILE_LENGTH
                    for tile in self.tiles:
                        tx = (tile["world_pos"][0] + self.TILE_LENGTH / 2) // self.TILE_LENGTH
                        ty = (tile["world_pos"][1] + self.TILE_LENGTH / 2) // self.TILE_LENGTH
                        if tx == cx and ty == cy:
                            if tile["type"] != "rescue":
                                self.scoreCan()
                            rel_pos_x = self.can_obj.position[0] - tile["world_pos"][0]
                            rel_pos_y = self.can_obj.position[1] - tile["world_pos"][1]
                            sx = 1
                            sy = 1
                            if rel_pos_x < 0:
                                rel_pos_x = -rel_pos_x
                                sx = -1
                            if rel_pos_y < 0:
                                rel_pos_y = -rel_pos_y
                                sy = -1
                            if rel_pos_y < self.TILE_LENGTH / 3 and rel_pos_x < self.TILE_LENGTH / 3:
                                break
                            under, right, under_right = self._checkConnectingRescueTiles(tile["world_pos"], sx, sy)
                            if under_right and under and right:
                                # We can't be outside of the area.
                                break
                            if under:
                                if rel_pos_x < self.TILE_LENGTH / 3:
                                    break
                            if right:
                                if rel_pos_y < self.TILE_LENGTH / 3:
                                    break
                            self.scoreCan()
                            break
                    else:
                        self.scoreCan()

        for i in range(len(self.tiles)):
            if self.tiles[i]["type"] == "follow":
                self.tiles[i]["checker"].tick(tick)
        # UI Tick
        if self._pressed:
            ScriptLoader.instance.object_map["controlsReset"].visual.image_path = "ui/controls_reset_pressed.png"
        else:
            ScriptLoader.instance.object_map["controlsReset"].visual.image_path = "ui/controls_reset_released.png"
        self.update_time()

    @stop_on_pause
    def update_time(self):
        self.time_tick += 1
        elapsed = datetime.timedelta(seconds=self.time_tick / ScriptLoader.instance.GAME_TICK_RATE)
        show = self.START_TIME - elapsed
        seconds = show.seconds
        minutes = seconds // 60
        seconds = seconds - minutes * 60
        ScriptLoader.instance.object_map["TimerText"].text = "{:02d}:{:02d}".format(minutes, seconds)

    def handleEvent(self, event):
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            m_pos = screenspace_to_worldspace(event.pos)
            shapes = World.instance.space.point_query(
                [float(v) for v in m_pos], 0.0, pymunk.ShapeFilter(mask=STATIC_CATEGORY)
            )
            for shape in shapes:
                if shape.shape.obj.key == "controlsReset":
                    self._pressed = True
                words = shape.shape.obj.key.split("-")
                if len(words) == 3 and words[0] == "Tile" and words[2] == "UI":
                    tile_index = int(words[1])
                    self.spawnAt(tile_index)
        if event.type == pygame.MOUSEBUTTONUP and event.button == 1:
            m_pos = screenspace_to_worldspace(event.pos)
            shapes = World.instance.space.point_query(
                [float(v) for v in m_pos], 0.0, pymunk.ShapeFilter(mask=STATIC_CATEGORY)
            )
            for shape in shapes:
                if (shape.shape.obj.key == "controlsReset") & self._pressed:
                    self.reset()
            self._pressed = False
        if event.type == pygame.MOUSEMOTION:
            m_pos = screenspace_to_worldspace(event.pos)
            shapes = World.instance.space.point_query(
                [float(v) for v in m_pos], 0.0, pymunk.ShapeFilter(mask=STATIC_CATEGORY)
            )
            for shape in shapes:
                words = shape.shape.obj.key.split("-")
                if len(words) == 3 and words[0] == "Tile" and words[2] == "UI":
                    tile_index = int(words[1])
                    # Highlight at tile.
                    if self.hover_rect.key not in ScreenObjectManager.instance.objects:
                        ScreenObjectManager.instance.registerVisual(self.hover_rect, self.hover_rect.key)
                    self.hover_rect.position = self.tiles[tile_index]["world_pos"]
                    break
            else:
                if self.hover_rect.key in ScreenObjectManager.instance.objects:
                    ScreenObjectManager.instance.unregisterVisual(self.hover_rect.key)

    TOUCH_PENALTY = 5
    MAX_TOUCH_PENALTY = 20

    def touchBot(self):
        self._touches += 1
        penalty = min(self.TOUCH_PENALTY, self.MAX_TOUCH_PENALTY - self._touch_points)
        self._touch_points += penalty
        self.decrementScore(penalty)
        self.touchesChanged()

    def touchesChanged(self):
        ScriptLoader.instance.object_map["touchesCount"].text = str(self._touches)
        ScriptLoader.instance.object_map["touchesScore"].text = str(-self._touch_points)

rescue_settings = {
    attr: ObjectSetting(RescueInteractor, attr)
    for attr in [
        "SHOW_FOLLOW_POINTS",
        "SHOW_ROBOT_COLLIDER",
        "ROBOT_CENTRE_RADIUS",
        "FOLLOW_POINT_RADIUS",
        "MAX_FOLLOW_DIST",
        "BOT_SPAWN_RADIUS",
        "BOT_SPAWN_ANGLE",
        "GAME_LENGTH_MINUTES",
        "BOT_SPAWN_POSITION",
        "CAN_SPAWN_POSITION",
        "TILE_DEFINITIONS",
    ]
}

from ev3sim.visual.settings.elements import NumberEntry, TextEntry, Checkbox, Button

def onClickMapEditor(filename):
    from ev3sim.visual.manager import ScreenObjectManager

    ScreenObjectManager.instance.pushScreen(
        ScreenObjectManager.SCREEN_RESCUE_EDIT,
        batch_file=find_abs(f"{filename}.sim", batch_locations()),
    )

class MapButton(Button):
    def resize(self, objects, index):
        super().resize(objects, index)
        if self.menu.creating:
            self.button.disable()

visual_settings = [
    {"height": lambda s: 90, "objects": [TextEntry("__filename__", "BATCH NAME", None, (lambda s: (0, 20)))]},
    {
        "height": (lambda s: 240 if s[0] < 540 else 140),
        "objects": [
            NumberEntry(["settings", "rescue", "GAME_LENGTH_MINUTES"], 5, "Time allowed", (lambda s: (0, 20)), float),
            NumberEntry(
                ["settings", "rescue", "ROBOT_CENTRE_RADIUS"],
                3,
                "Collision Radius",
                (lambda s: (0, 70) if s[0] < 540 else (s[0] / 2, 20)),
                float,
            ),
            NumberEntry(
                ["settings", "rescue", "MAX_FOLLOW_DIST"],
                8,
                "Lack of Progress",
                (lambda s: (0, 120) if s[0] < 540 else (0, 70)),
                float,
            ),
            NumberEntry(
                ["settings", "rescue", "FOLLOW_POINT_RADIUS"],
                1,
                "Follow Radius",
                (lambda s: (0, 170) if s[0] < 540 else (s[0] / 2, 70)),
                float,
            ),
        ],
    },
    {
        "height": (lambda s: 120 if s[0] < 540 else 70),
        "objects": [
            Checkbox(["settings", "rescue", "SHOW_FOLLOW_POINTS"], False, "Show Follow Points", (lambda s: (0, 20))),
            MapButton("Map Editor", (lambda s: (0, 70) if s[0] < 540 else (s[0] / 2, 20)), onClickMapEditor),
        ],
    },
]


import rlkit.torch.pytorch_util as ptu
from rlkit.data_management.env_replay_buffer import EnvReplayBuffer
from rlkit.envs.wrappers import NormalizedBoxEnv
from rlkit.launchers.launcher_util import setup_logger
from rlkit.samplers.data_collector import MdpPathCollector
from rlkit.torch.sac.policies import TanhGaussianPolicy, MakeDeterministic
from rlkit.torch.flowq.flowq2 import FlowQTrainer
from rlkit.torch.networks import FlattenMlp
from rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm

def experiment(variant):
    import gym
    expl_env = NormalizedBoxEnv(gym.make(args.exp_name+'-v2'))
    eval_env = NormalizedBoxEnv(gym.make(args.exp_name+'-v2'))
    obs_dim = expl_env.observation_space.low.size
    action_dim = eval_env.action_space.low.size

    M = variant['layer_size']
    vf1 = FlattenMlp(
        input_size=obs_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    vf2 = FlattenMlp(
        input_size=obs_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    alpha = FlattenMlp(
        input_size=obs_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    target_vf1 = FlattenMlp(
        input_size=obs_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    target_vf2 = FlattenMlp(
        input_size=obs_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    policy = TanhGaussianPolicy(
        obs_dim=obs_dim,
        action_dim=action_dim,
        hidden_sizes=[M, M],
        return_raw_action=True,
    )
    eval_policy = MakeDeterministic(policy)
    eval_path_collector = MdpPathCollector(
        eval_env,
        eval_policy,
    )
    expl_path_collector = MdpPathCollector(
        expl_env,
        policy,
    )
    replay_buffer = EnvReplayBuffer(
        variant['replay_buffer_size'],
        expl_env,
        store_raw_action=True,
    )
    trainer = FlowQTrainer(
        env=eval_env,
        policy=policy,
        vf1=vf1,
        vf2=vf2,
        alpha=alpha,
        target_vf1=target_vf1,
        target_vf2=target_vf2,
        **variant['trainer_kwargs']
    )
    algorithm = TorchBatchRLAlgorithm(
        trainer=trainer,
        exploration_env=expl_env,
        evaluation_env=eval_env,
        exploration_data_collector=expl_path_collector,
        evaluation_data_collector=eval_path_collector,
        replay_buffer=replay_buffer,
        **variant['algorithm_kwargs']
    )
    algorithm.to(ptu.device)
    algorithm.train()

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--exp_name', type=str, default='Hopper')
    parser.add_argument('--log_dir', type=str, default='FlowQ2')
    parser.add_argument('--lr', type=float, default=None)
    parser.add_argument('--cg', type=float, default=None) # clip gradient
    parser.add_argument('--sr', type=float, default=None) # reward scale
    parser.add_argument('--bs', type=int, default=None) # batch size
    parser.add_argument('--tui', type=int, default=None) # target update interval
    # parser.add_argument('--ae', type=int, default=None) # auto entropy, 0=False
    parser.add_argument('--epoch', type=int, default=None)
    parser.add_argument('--seed', type=int, default=0)
    parser.add_argument('--snapshot_mode', type=str, default="gap_and_last")
    parser.add_argument('--snapshot_gap', type=int, default=500)
    args = parser.parse_args()
    import os.path as osp
    pre_dir = './Data/'+args.exp_name
    main_dir = args.log_dir\
                +(('lr'+str(args.lr)) if args.lr else '')\
                +(('cg'+str(args.cg)) if args.cg else '')\
                +(('sr'+str(args.sr)) if args.sr else '')\
                +(('bs'+str(args.bs)) if args.bs else '')\
                +(('tui'+str(args.tui)) if args.tui else '')
                # +(('ae'+str(args.ae)) if args.ae==0 else '')
    log_dir = osp.join(pre_dir,main_dir,'seed'+str(args.seed))
    # noinspection PyTypeChecker
    variant = dict(
        algorithm="FlowQ",
        version="normal",
        layer_size=256,
        replay_buffer_size=int(1E6),
        algorithm_kwargs=dict(
            num_epochs=(args.epoch if args.epoch else 3000),
            num_eval_steps_per_epoch=5000,
            num_trains_per_train_loop=1000,
            num_expl_steps_per_train_loop=1000,
            min_num_steps_before_training=1000,
            max_path_length=1000,
            batch_size=(args.bs if args.bs else 256),
        ),
        trainer_kwargs=dict(
            discount=0.99,
            soft_target_tau=5e-3,
            target_update_period=(args.tui if args.tui else 1),
            policy_lr=(args.lr if args.lr else 3E-4),
            vf_lr=(args.lr if args.lr else 3E-4),
            reward_scale=(args.sr if args.sr else 1),
            # use_automatic_entropy_tuning=(False if args.ae==0 else True),
            clip_gradient=args.cg,
        ),
    )
    import os
    if not os.path.isdir(log_dir):
        os.makedirs(log_dir)
    with open(osp.join(log_dir,'variant.json'),'w') as out_json:
        import json
        json.dump(variant,out_json,indent=2)
    import sys
    cmd_input = 'python ' + ' '.join(sys.argv) + '\n'
    with open(osp.join(log_dir, 'cmd_input.txt'), 'a') as f:
        f.write(cmd_input)
    setup_logger(args.exp_name+'/'+main_dir, variant=variant,
                snapshot_mode=args.snapshot_mode, snapshot_gap=args.snapshot_gap,
                log_dir=log_dir)
    import numpy as np
    import torch
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    # ptu.set_gpu_mode(True)  # optionally set the GPU (default=False)
    experiment(variant)


'''
Class wrapper for imputation, transforming an input
Pandas data frame to an output Pandas data-frame
'''

import os
import sys
import os.path

import numpy as np
import scipy as sp
import pandas as pd

import circews.functions.forward_filling as bern_forward_fill
import circews.functions.util.array as mlhc_array
import circews.functions.util.math as mlhc_math

class TimegridderNoImpute:
    '''
    Function transforming the merged data representation of the ICU Bern data-set to
    a time-gridded version. Special version that does not impute, but puts NAN except at the
    exact grid interval where a measurement occured.
    '''
    def __init__(self, is_dim_reduced, grid_period=None, max_grid_length_days=None, use_adaptive_impute=True):

        self.is_dim_reduced=is_dim_reduced

        if grid_period is not None:
            self.grid_period=grid_period
        else:
            assert(False)

        if max_grid_length_days is not None:
            self.max_grid_length_days=max_grid_length_days
            self.max_grid_length_secs=self.max_grid_length_days*24*3600

        self.patient_id_key="PatientID"
        self.rel_datetime_key="RelDatetime"
        self.abs_datetime_key="AbsDatetime"

        # Default imputation parameters if the dictionary does not contain values
        self.default_fill_interval_secs=172800     # 2 DAYS (ARBITRARY)
        self.default_rolling_mean_secs=2*172800    # 4 DAYS (ARBITRARY)
        self.default_fill_interval_secs=172800     # 2 DAYS (ARBITRARY)

        if use_adaptive_impute is not None:
            self.use_adaptive_impute=use_adaptive_impute

        self.key_dict={"F": "female", "M": "male", "U": "unknown"}

    def set_static_table(self, df_static):
        ''' Set a static data table as context information in which we can lookup information
            about a patient'''
        self.df_static=df_static

    def set_typical_weight_dict(self, typical_weight_dict):
        ''' Set a typical weight dictionary as context information to use during information'''
        self.typical_weight_dict=typical_weight_dict

    def set_median_bmi_dict(self, median_bmi_dict):
        ''' Set a median BMI dictionary as context information to use during imputation'''
        self.median_bmi_dict=median_bmi_dict

    def set_var_encoding_map(self, var_encoding_map):
        ''' Set a variable encoding map that specifies for each of the variables in the data-frame the variable encoding'''
        self.var_encoding_map=var_encoding_map

    def set_imputation_params(self, global_impute_dict, interval_median_dict, interval_iqr_dict):
        ''' Set global imputation parameters of the transformer'''
        self.global_impute_dict=global_impute_dict
        self.interval_median_dict=interval_median_dict
        self.interval_iqr_dict=interval_iqr_dict

    def set_normal_vals(self, normal_val_map):
        ''' Set clinically normal values to be used for global replacement in the imputation algorithm'''
        self.normal_dict=normal_val_map

        # Convert to numeric encoding for some special variables
        if self.is_dim_reduced:
            self.normal_dict["vm57"]=0.0 # ECMO
            self.normal_dict["vm60"]=0.0 # ServoiMode
            self.normal_dict["vm66"]=0.0 # AirwayCode
            self.normal_dict["vm72"]=0.0 # Haemofiltration
            self.normal_dict["vm131"]="CUSTOM_FORMULA" # Weight
        else:
            self.normal_dict["v3845"]=0.0 # ServoiMode
            self.normal_dict["v15001552"]=0.0 # AirwayCode
            self.normal_dict["v10002508"]=0.0 # Haemofiltration
            self.normal_dict["v10000400"]="CUSTOM_FORMULA" # Weight

    def _check_state(self):
        ''' Assert correct initialization before the transformer is used'''
        assert(self.typical_weight_dict is not None)
        assert(self.df_static is not None)
        assert(self.median_bmi_dict is not None)
        assert(self.var_encoding_map is not None)
        assert(self.global_impute_dict is not None)
        assert(self.interval_median_dict is not None)
        assert(self.interval_iqr_dict is not None)
        assert(self.normal_dict is not None)

    def transform(self, patient_df, pid=None):
        ''' Transformer method, taking as input a data-frame with irregularly sampled input data. The method
            assumes that the data-frame contains a time-stamp column, and the data-frame is sorted along the first
            axis in non-decreasing order with respect to the timestamp column. Pass the <pid> of the patient stay
            as additional information'''
        self._check_state()
        static_table=self.df_static[self.df_static["PatientID"]==pid]

        # No static data, patient is not valid, exclude on-the-fly
        if static_table.shape[0]==0:
            print("WARNING: No static data in patient table...")
            return None

        # More than one row, select one of the rows arbitrarily
        if static_table.shape[0]>1:
            print("WARNING: More than one row in static table...")
            static_table=static_table.take([0],axis=0)

        static_height=float(static_table["Height"])
        static_gender=str(static_table["Sex"].values[0]).strip()
        assert(static_gender in ["F","M","U"])

        if static_gender in ["F", "M"]:
            typical_weight=self.typical_weight_dict[static_gender]
        else:
            typical_weight=(self.typical_weight_dict["M"]+self.typical_weight_dict["F"])/2.0

        personal_bmi=self.median_bmi_dict[self.key_dict[static_gender]]

        ## If either the endpoints or the features don't exist, log the failure but do nothing, the missing patients can be
        #  latter added as a new group to the output H5
        if patient_df.shape[0]==0:
            print("WARNING: p{} has missing features, skipping output generation...".format(pid))
            return None

        all_keys=list(set(patient_df.columns.values.tolist()).difference(set(["Datetime","PatientID","a_temp","m_pm_1", "m_pm_2"])))

        ts=patient_df["Datetime"]
        ts_arr=np.array(ts)
        n_ts=ts_arr.size

        if self.is_dim_reduced:
            hr=np.array(patient_df["vm1"])
        else:
            hr=np.array(patient_df["v200"])

        finite_hr=ts_arr[np.isfinite(hr)]

        if finite_hr.size==0:
            print("WARNING: Patient {} has no HR, ignoring patient...".format(pid))
            return None

        ts_min=ts_arr[np.isfinite(hr)][0]
        ts_max=ts_arr[np.isfinite(hr)][-1]
        max_ts_diff=(ts_max-ts_min)/np.timedelta64(1,'s')

        time_grid=np.arange(0.0,min(max_ts_diff+1.0,self.max_grid_length_secs),self.grid_period)
        time_grid_abs=[ts_min+pd.Timedelta(seconds=time_grid[idx]) for idx in range(time_grid.size)]
        imputed_df_dict={}
        imputed_df_dict[self.patient_id_key]=[int(pid)]*time_grid.size
        imputed_df_dict[self.rel_datetime_key]=time_grid
        imputed_df_dict[self.abs_datetime_key]=time_grid_abs

        ## There is nothing to do if the patient has no records, just return...
        if n_ts==0:
            print("WARNING: p{} has an empty record, skipping output generation...".format(patient))
            return None

        ## Initialize the storage for the imputed time grid, NANs for the non-pharma, 0 for pharma.
        for col in all_keys:
            if col[0]=="p":
                imputed_df_dict[col]=np.zeros(time_grid.size)
            elif col[0]=="v":
                imputed_df_dict[col]=mlhc_array.empty_nan(time_grid.size)
            else:
                print("ERROR: Invalid variable type")
                assert(False)

        imputed_df=pd.DataFrame(imputed_df_dict)
        norm_ts=np.array(ts-ts_min)/np.timedelta64(1,'s')

        # Schedule for order of variable imputation
        if self.is_dim_reduced:
            all_keys.remove("vm131")
            all_keys=["vm131"]+all_keys
        else:
            all_keys.remove("v10000400")
            all_keys=["v10000400"]+all_keys

        ## Impute all variables independently, with the two relevant cases pharma variable and other variable,
        #  distinguishable from the variable prefix. We enforce that weight is the first variable to be imputed, so that
        #  its time-gridded information can later be used by other custom formulae imputations that depend on it.
        for var_idx,variable in enumerate(all_keys):
            df_var=patient_df[variable]
            assert(n_ts==df_var.shape[0]==norm_ts.size)

            ## Non-pharma variable case
            if variable[0]=="v":
                valid_normal=False
                var_encoding=self.var_encoding_map[variable]

                # Saved a value in the dict of normal values
                if variable in self.normal_dict:
                    saved_normal_var=self.normal_dict[variable]

                    # Saved normal value is already numeric, no need to encode it here...
                    if mlhc_math.is_numeric(saved_normal_var) and np.isfinite(saved_normal_var):
                        global_impute_val=saved_normal_var
                        valid_normal=True

                # Could not determine a valid normal value, have to fall back to pre-computed global statistic
                if not valid_normal:

                    # Fill in the weight using BMI calculations
                    if variable in ["v10000400","vm131"]:

                        # If we have an observed height can use BMI
                        if np.isfinite(static_height):
                            global_impute_val=personal_bmi*(static_height/100)**2
                        else:
                            global_impute_val=typical_weight

                    # Fill in with the global statistic
                    elif variable in self.global_impute_dict:
                        global_impute_val=self.global_impute_dict[variable]

                    # Rare case, no observation in the imputation data-set
                    else:
                        global_impute_val=np.nan

                # Default values where median/IQR interval not saved
                if variable not in self.interval_median_dict:
                    fill_interval_secs=self.default_fill_interval_secs
                    rolling_mean_secs=self.default_rolling_mean_secs
                    fill_interval_secs=self.default_fill_interval_secs

                # We have to impose minimum period to have boundary conditions where the backward window for
                # slope estimation is empty or an observation is not even filled to the next grid point to the right.
                else:
                    med_interval=self.interval_median_dict[variable]
                    iqr_interval=self.interval_iqr_dict[variable]
                    base_val=med_interval+2*iqr_interval
                    fill_interval_secs=max(self.grid_period, base_val)
                    rolling_mean_secs=max(2*self.grid_period, 2*base_val)
                    return_mean_secs=max(2*self.grid_period, base_val)

                raw_col=np.array(df_var)
                assert(raw_col.size==norm_ts.size)
                observ_idx=np.isfinite(raw_col)
                observ_ts=norm_ts[observ_idx]
                observ_val=raw_col[observ_idx]

                ## No values have been observed for this variable, it has to be imputed using the global mean
                if observ_val.size==0:
                    est_vals=mlhc_array.value_empty(time_grid.size,np.nan)
                    imputed_df[variable]=est_vals
                    imputed_df["{}_IMPUTED_STATUS_CUM_COUNT".format(variable)]=np.zeros(time_grid.size)
                    imputed_df["{}_IMPUTED_STATUS_TIME_TO".format(variable)]=mlhc_array.value_empty(time_grid.size,-1.0)
                    continue

                assert(np.isfinite(observ_val).all())
                assert(np.isfinite(observ_ts).all())

                if self.use_adaptive_impute:

                    # Formulae imputation
                    if variable in ["v1000","v1010","v10020000","v30005010","v30005110","vm13","vm24","vm31","vm32"]:
                        existing_weight_col=np.array(imputed_df["vm131"]) if self.is_dim_reduced else np.array(imputed_df["v10000400"])
                        est_vals,cum_count_ts,time_to_last_ms=bern_forward_fill.impute_forward_fill_new_no_impute(observ_ts,observ_val,time_grid, global_impute_val, self.grid_period,
                                                                                fill_interval_secs= fill_interval_secs, rolling_mean_secs= rolling_mean_secs,
                                                                                return_mean_secs= return_mean_secs, var_type="non_pharma", var_encoding=var_encoding,
                                                                                variable_id=variable, weight_imputed_col=existing_weight_col, static_height=static_height,
                                                                                personal_bmi=personal_bmi)
                    elif variable in ["v10000400","vm131"]:
                        est_vals,cum_count_ts,time_to_last_ms=bern_forward_fill.impute_forward_fill_new_no_impute(observ_ts, observ_val, time_grid, global_impute_val, self.grid_period,var_type="weight")
                    else:
                        est_vals,cum_count_ts,time_to_last_ms=bern_forward_fill.impute_forward_fill_new_no_impute(observ_ts,observ_val,time_grid, global_impute_val, self.grid_period,
                                                                                fill_interval_secs= fill_interval_secs, rolling_mean_secs= rolling_mean_secs,
                                                                                return_mean_secs= return_mean_secs, var_type="non_pharma", var_encoding=var_encoding,
                                                                                variable_id=variable)

                else:
                    assert(False)
                    est_vals=bern_forward_fill.impute_forward_fill(observ_ts, observ_val, time_grid, global_mean_var)

                imputed_df[variable]=est_vals
                imputed_df["{}_IMPUTED_STATUS_CUM_COUNT".format(variable)]=cum_count_ts
                imputed_df["{}_IMPUTED_STATUS_TIME_TO".format(variable)]=time_to_last_ms

            ## Pharma variable case, the doses have to be recomputed to the time-grid. The global imputation value is 0, because the rate assumed w/o observation
            #  is 0 (no medication flow)
            elif variable[0]=="p":
                global_impute_val=0.0
                raw_col=np.array(df_var)
                assert(raw_col.size==norm_ts.size)
                observ_idx=np.isfinite(raw_col)
                observ_ts=norm_ts[observ_idx]
                observ_val=raw_col[observ_idx]

                ## No values have been observed for this pharma-variable, leave Zero in this series
                if observ_val.size==0:
                    continue

                assert(np.isfinite(observ_val).all())
                assert(np.isfinite(observ_ts).all())
                est_vals,cum_count_ts,time_to_last_ms=bern_forward_fill.impute_forward_fill_new_no_impute(observ_ts, observ_val, time_grid, global_impute_val, self.grid_period,var_type="pharma")
                imputed_df[variable]=est_vals
            else:
                print("ERROR: Invalid variable, exiting...")
                assert(False)

        return imputed_df


# --------------------------------------------------------
# Pytorch Multi-GPU Faster R-CNN
# Licensed under The MIT License [see LICENSE for details]
# Written by Jiasen Lu, Jianwei Yang, based on code from Ross Girshick
# --------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import _init_paths
import os
import sys
import numpy as np
import argparse
import pprint
import pdb
import time

import cv2

import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.optim as optim
import pickle
from roi_data_layer.roidb import combined_roidb
from roi_data_layer.roibatchLoader import roibatchLoader
from model.utils.config import cfg, cfg_from_file, cfg_from_list, get_output_dir
from model.rpn.bbox_transform import clip_boxes
# from model.nms.nms_wrapper import nms
from model.roi_layers import nms
from model.rpn.bbox_transform import bbox_transform_inv
from model.utils.net_utils import save_net, load_net, vis_detections
from model.faster_rcnn.vgg16 import vgg16
from model.faster_rcnn.resnet import resnet
from PIL import Image
from torchvision.utils import save_image
import cv2
from torch.utils.data import DataLoader
from srgan_datasets import *
from srgan import *
import torch.nn.functional as F
from datasets.voc_eval import parse_rec
try:
    xrange          # Python 2
except NameError:
    xrange = range  # Python 3

classes = ('__background__',  # always index 0
           'A.bes(H)','A.bes(T)','A.bes','A.bic(H)','A.bic(T)','A.bic',
           'A.fuj(H)','A.fuj(T)','A.fuj','B.xyl(H)','B.xyl(T)','B.xyl',
           'C.ele(H)','C.ele(T)','C.ele','M.ent(H)','M.ent(T)','M.ent',
           'M.gra(H)','M.gra(T)','M.gra','M.inc(H)','M.inc(T)','M.inc',
           'P.cof(H)','P.cof(T)','P.cof','P.vul(H)','P.vul(T)','P.vul',
           'P.spe(H)','P.spe(T)','P.spe','H.sp(H)','H.sp(T)','H.sp',
           'M.ams(H)' ,'M.ams(T)','M.ams'
           )###################
def parse_args():
  """
  Parse input arguments
  """
  parser = argparse.ArgumentParser(description='Train a Fast R-CNN network')
  parser.add_argument('--dataset', dest='dataset',
                      help='training dataset',
                      default='pascal_voc', type=str)
  parser.add_argument('--cfg', dest='cfg_file',
                      help='optional config file',
                      default='cfgs/res101.yml', type=str)
  parser.add_argument('--net', dest='net',
                      help='vgg16, res50, res101, res152',
                      default='res101', type=str)
  parser.add_argument('--set', dest='set_cfgs',
                      help='set config keys', default=None,
                      nargs=argparse.REMAINDER)
  parser.add_argument('--load_dir', dest='load_dir',
                      help='directory to load models', default="models",
                      type=str)
  parser.add_argument('--cuda', dest='cuda',
                      help='whether use CUDA',
                      action='store_true')
  parser.add_argument('--ls', dest='large_scale',
                      help='whether use large imag scale',
                      action='store_true')
  parser.add_argument('--mGPUs', dest='mGPUs',
                      help='whether use multiple GPUs',
                      action='store_true')
  parser.add_argument('--cag', dest='class_agnostic',
                      help='whether perform class_agnostic bbox regression',
                      action='store_true')
  parser.add_argument('--parallel_type', dest='parallel_type',
                      help='which part of model to parallel, 0: all, 1: model before roi pooling',
                      default=0, type=int)
  parser.add_argument('--checksession', dest='checksession',
                      help='checksession to load model',
                      default=1, type=int)
  parser.add_argument('--checkepoch', dest='checkepoch',
                      help='checkepoch to load network',
                      default=1, type=int)
  parser.add_argument('--checkpoint', dest='checkpoint',
                      help='checkpoint to load network',
                      default=10021, type=int)
  parser.add_argument('--vis', dest='vis',
                      help='visualization mode',
                      action='store_true')
  args = parser.parse_args()
  return args

def parse_args_for_srgan():
  os.makedirs("srgan/images", exist_ok=True)
  os.makedirs("srgan/saved_models", exist_ok=True)

  parser = argparse.ArgumentParser()
  parser.add_argument("--epoch", type=int, default=500 , help="epoch to start training from")
  parser.add_argument("--n_epochs", type=int, default=501, help="number of epochs of training")
  parser.add_argument("--dataset_name", type=str, default="img_align_celeba", help="name of the dataset")
  parser.add_argument("--batch_size", type=int, default=1, help="size of the batches")
  parser.add_argument("--lr", type=float, default=0.00001, help="adam: learning rate")
  parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
  parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of second order momentum of gradient")
  parser.add_argument("--decay_epoch", type=int, default=100, help="epoch from which to start lr decay")
  #parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
  parser.add_argument("--hr_height", type=int, default=1024, help="high res. image height")
  parser.add_argument("--hr_width", type=int, default=1024, help="high res. image width")
  parser.add_argument("--channels", type=int, default=3, help="number of image channels")
  parser.add_argument("--sample_interval", type=int, default=50, help="interval between saving image samples")
  parser.add_argument("--checkpoint_interval", type=int, default=100, help="interval between model checkpoints")

  opt = parser.parse_args([])

  return opt

lr = cfg.TRAIN.LEARNING_RATE
momentum = cfg.TRAIN.MOMENTUM
weight_decay = cfg.TRAIN.WEIGHT_DECAY

def load_gt_box(annopath,
             imagesetfile,
             classname,
             cachedir):
  if not os.path.isdir(cachedir):
      os.mkdir(cachedir)
  cachefile = os.path.join(cachedir, '%s_annots.pkl' % imagesetfile)
  # read list of images
  with open(imagesetfile, 'r') as f:
    lines = f.readlines()
  imagenames = [x.strip() for x in lines]

  if not os.path.isfile(cachefile):
    # load annotations
    recs = {}
    for i, imagename in enumerate(imagenames):
      recs[imagename] = parse_rec(annopath.format(imagename))
      if i % 100 == 0:
        print('Reading annotation for {:d}/{:d}'.format(
          i + 1, len(imagenames)))
    # save
    print('Saving cached annotations to {:s}'.format(cachefile))
    with open(cachefile, 'wb') as f:
      pickle.dump(recs, f)
  else:
    # load
    with open(cachefile, 'rb') as f:
      try:
        recs = pickle.load(f)
      except:
        recs = pickle.load(f, encoding='bytes')
  #print(recs)
  # extract gt objects for this class
  class_recs = {}
  npos = 0
  for imagename in imagenames:
    R = [obj for obj in recs[imagename] if obj['name'] == classname]
    bbox = np.array([x['bbox'] for x in R])
    difficult = np.array([x['difficult'] for x in R]).astype(np.bool)
    det = [False] * len(R)
    npos = npos + sum(~difficult)
    class_recs[imagename] = {'bbox': bbox,
                             'difficult': difficult,
                             'det': det}
  #print(class_recs)
  #print(len(class_recs))
  return class_recs

def iou(bb1, bb2):#########################
  """ check if overlap"""
  #assert bb1[0] < bb1[2]
  #assert bb1[1] < bb1[3]
  #assert bb2[0] < bb2[2]
  #assert bb2[1] < bb2[3]

  # determine the coordinates of the intersection rectangle
  #print(bb1[0], bb2[0])
  x_left = max(bb1[0], bb2[0])
  y_top = max(bb1[1], bb2[1])
  x_right = min(bb1[2], bb2[2])
  y_bottom = min(bb1[3], bb2[3])
  iw = x_right - x_left
  ih = y_bottom - y_top
  inters = iw * ih
  # union
  uni = ((bb1[2]-bb1[0])*(bb1[3]-bb1[1]) + (bb2[2]-bb2[0])*(bb2[3]-bb2[1]) - inters)

  overlaps = inters / uni
  return overlaps

def Area(vertex):
  width = vertex[2] - vertex[0]
  height = vertex[3] - vertex[1]
  area = width*height
  return area

if __name__ == '__main__':

  args = parse_args()
  args_sr = parse_args_for_srgan()
  print('Called with args:')
  print(args)

  if torch.cuda.is_available() and not args.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with --cuda")

  np.random.seed(cfg.RNG_SEED)
  if args.dataset == "pascal_voc":
      args.imdb_name = "voc_2007_trainval"
      args.imdbval_name = "voc_2007_test"
      args.set_cfgs = ['ANCHOR_SCALES', '[8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]']
  elif args.dataset == "pascal_voc_0712":
      args.imdb_name = "voc_2007_trainval+voc_2012_trainval"
      args.imdbval_name = "voc_2007_test"
      args.set_cfgs = ['ANCHOR_SCALES', '[8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]']
  elif args.dataset == "coco":
      args.imdb_name = "coco_2014_train+coco_2014_valminusminival"
      args.imdbval_name = "coco_2014_minival"
      args.set_cfgs = ['ANCHOR_SCALES', '[4, 8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]']
  elif args.dataset == "imagenet":
      args.imdb_name = "imagenet_train"
      args.imdbval_name = "imagenet_val"
      args.set_cfgs = ['ANCHOR_SCALES', '[8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]']
  elif args.dataset == "vg":
      args.imdb_name = "vg_150-50-50_minitrain"
      args.imdbval_name = "vg_150-50-50_minival"
      args.set_cfgs = ['ANCHOR_SCALES', '[4, 8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]']

  args.cfg_file = "cfgs/{}_ls.yml".format(args.net) if args.large_scale else "cfgs/{}.yml".format(args.net)

  if args.cfg_file is not None:
    cfg_from_file(args.cfg_file)
  if args.set_cfgs is not None:
    cfg_from_list(args.set_cfgs)

  print('Using config:')
  pprint.pprint(cfg)

  cfg.TRAIN.USE_FLIPPED = False
  imdb, roidb, ratio_list, ratio_index = combined_roidb(args.imdbval_name, False)
  imdb.competition_mode(on=True)

  print('{:d} roidb entries'.format(len(roidb)))

  input_dir = args.load_dir + "/" + args.net + "/" + args.dataset
  if not os.path.exists(input_dir):
    raise Exception('There is no input directory for loading network from ' + input_dir)
  load_name = os.path.join(input_dir,
    'faster_rcnn_{}_{}_{}.pth'.format(args.checksession, args.checkepoch, args.checkpoint))

  # initilize the network here.
  if args.net == 'vgg16':
    fasterRCNN = vgg16(imdb.classes, pretrained=False, class_agnostic=args.class_agnostic)
  elif args.net == 'res101':
    fasterRCNN = resnet(imdb.classes, 101, pretrained=False, class_agnostic=args.class_agnostic)
  elif args.net == 'res50':
    fasterRCNN = resnet(imdb.classes, 50, pretrained=False, class_agnostic=args.class_agnostic)
  elif args.net == 'res152':
    fasterRCNN = resnet(imdb.classes, 152, pretrained=False, class_agnostic=args.class_agnostic)
  else:
    print("network is not defined")
    pdb.set_trace()

  fasterRCNN.create_architecture()

  print("load checkpoint %s" % (load_name))
  checkpoint = torch.load(load_name)
  fasterRCNN.load_state_dict(checkpoint['model'])
  if 'pooling_mode' in checkpoint.keys():
    cfg.POOLING_MODE = checkpoint['pooling_mode']

  print('load model successfully!')
  # initilize the tensor holder here.
  im_data = torch.FloatTensor(1)
  im_info = torch.FloatTensor(1)
  num_boxes = torch.LongTensor(1)
  gt_boxes = torch.FloatTensor(1)

  # ship to cuda
  if args.cuda:
    im_data = im_data.cuda()
    im_info = im_info.cuda()
    num_boxes = num_boxes.cuda()
    gt_boxes = gt_boxes.cuda()

  # make variable
  im_data = Variable(im_data)
  im_info = Variable(im_info)
  num_boxes = Variable(num_boxes)
  gt_boxes = Variable(gt_boxes)

  if args.cuda:
    cfg.CUDA = True

  if args.cuda:
    fasterRCNN.cuda()

  start = time.time()
  max_per_image = 100

  vis = args.vis

  if vis:
    thresh = 0.0
  else:
    thresh = 0.0

  save_name = 'faster_rcnn_10'
  num_images = len(imdb.image_index)

  all_boxes = [[[] for _ in range(num_images)]
               for _ in range(imdb.num_classes)]
  output_dir = get_output_dir(imdb, save_name)
  dataset = roibatchLoader(roidb, ratio_list, ratio_index, 1, \
                        imdb.num_classes, training=False, normalize = False)
  dataloader = torch.utils.data.DataLoader(dataset, batch_size=1,
                            shuffle=False, num_workers=0,
                            pin_memory=True)

  data_iter = iter(dataloader)

  _t = {'im_detect': time.time(), 'misc': time.time()}
  det_file = os.path.join(output_dir, 'detections.pkl')

  fasterRCNN.eval()
  empty_array = np.transpose(np.array([[],[],[],[],[]]), (1,0))
  for i in range(num_images):

      data = next(data_iter)
      with torch.no_grad():
              im_data.resize_(data[0].size()).copy_(data[0])
              im_info.resize_(data[1].size()).copy_(data[1])
              gt_boxes.resize_(data[2].size()).copy_(data[2])
              num_boxes.resize_(data[3].size()).copy_(data[3])

      #print(im_data.shape)
      #print(im_info.shape)
      #print(gt_boxes)
      #print(num_boxes)
      det_tic = time.time()
      rois, cls_prob, bbox_pred, \
      rpn_loss_cls, rpn_loss_box, \
      RCNN_loss_cls, RCNN_loss_bbox, \
      rois_label = fasterRCNN(im_data, im_info, gt_boxes, num_boxes)

      scores = cls_prob.data
      boxes = rois.data[:, :, 1:5]
      if cfg.TEST.BBOX_REG:
          # Apply bounding-box regression deltas
          box_deltas = bbox_pred.data
          if cfg.TRAIN.BBOX_NORMALIZE_TARGETS_PRECOMPUTED:
          # Optionally normalize targets by a precomputed mean and stdev
            if args.class_agnostic:
                box_deltas = box_deltas.view(-1, 4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() \
                           + torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda()
                box_deltas = box_deltas.view(1, -1, 4)
            else:
                box_deltas = box_deltas.view(-1, 4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() \
                           + torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda()
                box_deltas = box_deltas.view(1, -1, 4 * len(imdb.classes))

          pred_boxes = bbox_transform_inv(boxes, box_deltas, 1)
          pred_boxes = clip_boxes(pred_boxes, im_info.data, 1)
      else:
          # Simply repeat the boxes, once for each class
          pred_boxes = np.tile(boxes, (1, scores.shape[1]))

      pred_boxes /= data[1][0][2].item()

      scores = scores.squeeze()
      pred_boxes = pred_boxes.squeeze()
      det_toc = time.time()
      detect_time = det_toc - det_tic
      misc_tic = time.time()

      #print(scores[:,1:3].shape)
      #print(pred_boxes[:,4:12].shape)
      ##############################   decline head-tail overlapping
      new_pred_boxes = torch.cuda.FloatTensor(300, 160).zero_()
      new_scores = torch.cuda.FloatTensor(300,40).zero_()
      for k in range(13):
        b = torch.cat((pred_boxes[:,12*k+4:12*k+8],pred_boxes[:,12*k+8:12*k+12]),0)
        s = torch.cat((scores[:,3*k+1],scores[:,3*k+2]),0)
        keep = nms(b, s, 0.2)
        #new head class
        idx = [g for g in range(len(keep)) if keep[g] <300]
        new_pred_boxes[:len(keep[idx]),12*k+4:12*k+8] = b[keep[idx]]
        new_scores[:len(keep[idx]),3*k+1] = s[keep[idx]]
        #new tail class
        idx = [g for g in range(len(keep)) if keep[g] >=300]
        new_pred_boxes[:len(keep[idx]),12*k+8:12*k+12] = b[keep[idx]]
        new_scores[:len(keep[idx]),3*k+2] = s[keep[idx]]
        #new full length class = original
        new_pred_boxes[:,12*k+12:12*k+16] = pred_boxes[:,12*k+12:12*k+16]
        new_scores[:,3*k+3] = scores[:,3*k+3]
      if vis:
          im = cv2.imread(imdb.image_path_at(i))
          im2show = np.copy(im)

      for j in range(1, imdb.num_classes):
          inds = torch.nonzero(new_scores[:,j]>thresh).view(-1)
          # if there is det
          if inds.numel() > 0:
            cls_scores = new_scores[:,j][inds]
            _, order = torch.sort(cls_scores, 0, True)
            if args.class_agnostic:
              cls_boxes = new_pred_boxes[inds, :]
            else:
              cls_boxes = new_pred_boxes[inds][:, j * 4:(j + 1) * 4]
            #print(cls_boxes.shape)
            #print(cls_scores.unsqueeze(1).shape)
            cls_dets = torch.cat((cls_boxes, cls_scores.unsqueeze(1)), 1)
            # cls_dets = torch.cat((cls_boxes, cls_scores), 1)
            cls_dets = cls_dets[order]
            keep = nms(cls_boxes[order, :], cls_scores[order], cfg.TEST.NMS)
            cls_dets = cls_dets[keep.view(-1).long()]

            if vis:
              im2show = vis_detections(im2show, imdb.classes[j], cls_dets.cpu().numpy(), 0.3)
            all_boxes[j][i] = cls_dets.cpu().numpy()
          else:
            all_boxes[j][i] = empty_array
      #print(exist_classes)
      #for k, j in enumerate(exist_classes):
      #  all_boxes[j][i] = exist_dets[k]
      #print(all_boxes)

      # Limit to max_per_image detections *over all classes*
      if max_per_image > 0:
          #print(all_boxes[3][i][:,-1])
          image_scores = np.hstack([all_boxes[j][i][:,-1]
                                    for j in range(1, imdb.num_classes)])
          if len(image_scores) > max_per_image:
              image_thresh = np.sort(image_scores)[-max_per_image]
              for j in xrange(1, imdb.num_classes):
                  keep = np.where(all_boxes[j][i][:, -1] >= image_thresh)[0]
                  all_boxes[j][i] = all_boxes[j][i][keep, :]

      misc_toc = time.time()
      nms_time = misc_toc - misc_tic

      sys.stdout.write('im_detect: {:d}/{:d} {:.3f}s {:.3f}s   \r' \
          .format(i + 1, num_images, detect_time, nms_time))
      sys.stdout.flush()

      if vis:
          cv2.imwrite('result.png', im2show)
          pdb.set_trace()
          #cv2.imshow('test', im2show)
          #cv2.waitKey(0)
  #print(all_boxes[1][0][0])
  print(torch.cuda.current_device())
  with torch.cuda.device(torch.cuda.current_device()):
    torch.cuda.empty_cache()
  #################################### filter imgs need to do SRGAN-preprocessing
  annopath = '/home/jason/faster-rcnn.pytorch-1.0/data/VOCdevkit2007/VOC2007/Annotations/{:s}.xml'
  imagesetfile = '/home/jason/faster-rcnn.pytorch-1.0/data/VOCdevkit2007/VOC2007/ImageSets/Main/test.txt'
  cachedir = '/home/jason/faster-rcnn.pytorch-1.0/data/VOCdevkit2007/annotations_cache'
  image_file = '/home/jason/faster-rcnn.pytorch-1.0/data/VOCdevkit2007/VOC2007/JPEGImages'
  f = open(imagesetfile)
  new_indexes = []
  img_ids = []
  new_gt_boxes = []
  for line in f:
    img_ids.append(line.splitlines())
  img_ids = np.squeeze(img_ids)
  for i in range(num_images):
    for j in range(1, imdb.num_classes):
      gt_boxes_1 = load_gt_box(annopath,imagesetfile,classes[j],cachedir)
      if not np.any(all_boxes[j][i]):
        continue

      if len(gt_boxes_1[img_ids[i]]['bbox']) == 0:
        continue
      else:# 1 GT box in single image for a single class
        gt_b = gt_boxes_1[img_ids[i]]['bbox']
        #print(gt_b)
        z = 0
        for m in range(len(all_boxes[j][i])):
          for n in range(len(gt_b)):
            det_b = [int(l) for l in all_boxes[j][i][m][:4]]
            #print(all_boxes[j][i][m][4], iou(det_b, gt_b[n]), imdb.image_index[j])
            if all_boxes[j][i][m][4] > 0.5 and all_boxes[j][i][m][4] < 0.8 \
                and iou(det_b, gt_b[n]) > 0.5 and classes[j][-1]==")":
              print("srgan beginning......")
              new_indexes.append(img_ids[i]+"_"+classes[j]+"_"+str(z))
              print(len(new_indexes))#, all_boxes[j][i][m][4], iou(det_b, gt_b[n]))
              img_path = os.path.join(image_file, img_ids[i]+".JPG")
              img = Image.open(img_path)
              img = np.asarray(img)
              quaterx = int(img.shape[1]*1/4)
              quatery = int(img.shape[0]*1/4)
              x1_padding = 0
              y1_padding = 0
              x2_padding = 0
              y2_padding = 0
              print(img.shape)
              if Area(det_b) >= Area(gt_b[n]):
                x1, y1, x2, y2 = det_b
                print("det_b: " + str(det_b))
                if x1 > quaterx:
                  x1-=quaterx
                  x1_padding = quaterx
                else:
                  x1 = 0
                  x1_padding = x1
                if x2 < img.shape[0]-quaterx:
                  x2+= quaterx
                  x2_padding = quaterx
                else:
                  x2 = img.shape[0]-1
                  x2_padding = img.shape[0] - x2-1
                if y1 > quatery:
                  y1 -=quatery
                  y1_padding = quatery
                else:
                  y1 = 0
                  y1_padding = y1
                if y2 < img.shape[1]-quatery:
                  y2+=quatery
                  y2_padding = quatery
                else:
                  y2= img.shape[1]-1
                  y2_padding = img.shape[1] - y2-1
              else:
                x1, y1, x2, y2 = gt_b[n]
                print("gt_b: "+str(gt_b))
                if x1 > quaterx:
                  x1-=quaterx
                  x1_padding = quaterx
                else:
                  x1 = 0
                  x1_padding = x1
                if x2 < img.shape[0]-quaterx:
                  x2+= quaterx
                  x2_padding = quaterx
                else:
                  x2 = img.shape[0]-1
                  x2_padding = img.shape[0] - x2-1
                if y1 > quatery:
                  y1 -=quatery
                  y1_padding = quatery
                else:
                  y1 = 0
                  y1_padding = y1
                if y2 < img.shape[1]-quatery:
                  y2+=quatery
                  y2_padding = quatery
                else:
                  y2= img.shape[1]-1
                  y2_padding = img.shape[1] - y2-1
              x1, y1, x2, y2= int(x1),int(y1),int(x2), int(y2)

              new_gt_boxes.append([x1_padding, y1_padding, x2-x1-x1_padding-x2_padding, \
                                   y2-y1-y1_padding-y2_padding])# whole photo
              srgan_in = img[y1:y2 ,x1:x2 ,:]
              srgan_in = srgan_in[...,::-1]#rgb->bgr
              print(x1,y1,x2,y2,srgan_in.shape)
              cv2.imwrite(os.path.join("srgan/srgan_input", img_ids[i]+"_"+classes[j]+"_"+str(z)+".JPG"), srgan_in)
              print("save input: %s" %(img_ids[i]+"_"+classes[j]+"_"+str(z)))
              z+=1
              all_boxes[j][i][m] = np.append(gt_b[n], 1.0)# turn original pred box to gt box

  with torch.cuda.device(torch.cuda.current_device()):
    torch.cuda.empty_cache()
  dataloader = DataLoader(
    ImageDataset("srgan/srgan_input", hr_shape=(1024,1024)),
    batch_size=1,
    shuffle=True,
    num_workers=0,
  )
  #gan_output = srgan(args_sr, dataloader)
  srgan(args_sr, dataloader)
  #print("length of data: %d"%len(gan_output))
  print("srgan finish......")
  with torch.cuda.device(torch.cuda.current_device()):
    torch.cuda.empty_cache()
  # re-test srgan output
  dataloader1 = DataLoader(
    ImageDataset("srgan/srgan_output", hr_shape=(1024,1024)),
    batch_size=1,
    shuffle=True,
    num_workers=0,
  )
  all_boxes_1 = [[[] for _ in range(len(dataloader1))]
               for _ in range(imdb.num_classes)]

  for i, gan_img in enumerate(dataloader1):
  #for i in range(len(dataloader1)):
    #gan_img = gan_output[i]
    #print(gan_img)
    arr = np.append(gan_img["origin_size"][0][0].numpy(), gan_img["origin_size"][1][0].numpy())
    gan_img_os = F.interpolate(gan_img['hr'], size=(arr[0],arr[1]), mode='bilinear')
    r = 600 / gan_img_os.shape[2]
    gan_info = np.array([[gan_img_os.shape[2], gan_img_os.shape[3], r]])
    with torch.no_grad():
      gan_img_600 = F.interpolate(gan_img_os, scale_factor=r, mode="bilinear").cuda()
      gan_info = torch.from_numpy(gan_info).cuda()
      gt_boxes
      num_boxes
    #print(gan_img.shape)
    #print(gan_info.shape)
    #print(gt_boxes)
    #print(num_boxes)
    det_tic = time.time()
    rois_1, cls_prob_1, bbox_pred_1, \
    rpn_loss_cls_1, rpn_loss_box_1, \
    RCNN_loss_cls_1, RCNN_loss_bbox_1, \
    rois_label_1 = fasterRCNN(gan_img_600, gan_info, gt_boxes, num_boxes)

    scores_1 = cls_prob_1.data
    boxes_1 = rois_1.data[:, :, 1:5]
    #print(data)
    if cfg.TEST.BBOX_REG:
        # Apply bounding-box regression deltas
        box_deltas = bbox_pred_1.data
        if cfg.TRAIN.BBOX_NORMALIZE_TARGETS_PRECOMPUTED:
        # Optionally normalize targets by a precomputed mean and stdev
          if args.class_agnostic:
              box_deltas_1 = box_deltas.view(-1, 4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() \
                         + torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda()
              box_deltas_1 = box_deltas.view(1, -1, 4)
          else:
              box_deltas_1 = box_deltas.view(-1, 4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() \
                         + torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda()
              box_deltas_1 = box_deltas.view(1, -1, 4 * len(imdb.classes))

        pred_boxes_1 = bbox_transform_inv(boxes, box_deltas_1, 1)
        pred_boxes_1 = clip_boxes(pred_boxes_1, gan_info.data, 1)
    else:
        # Simply repeat the boxes, once for each class
        pred_boxes_1 = np.tile(boxes_1, (1, scores.shape[1]))

    pred_boxes_1 /= data[1][0][2].item()

    scores_1 = scores_1.squeeze()
    pred_boxes_1 = pred_boxes_1.squeeze()
    det_toc = time.time()
    detect_time = det_toc - det_tic
    misc_tic = time.time()

    ##############################   decline head-tail overlapping
    new_pred_boxes = torch.cuda.FloatTensor(300, 160).zero_()
    new_scores = torch.cuda.FloatTensor(300,40).zero_()
    for k in range(13):
      b = torch.cat((pred_boxes_1[:,12*k+4:12*k+8],pred_boxes_1[:,12*k+8:12*k+12]),0)
      s = torch.cat((scores_1[:,3*k+1],scores_1[:,3*k+2]),0)
      keep = nms(b, s, 0.2)
      #new head class
      idx = [g for g in range(len(keep)) if keep[g] <300]
      new_pred_boxes[:len(keep[idx]),12*k+4:12*k+8] = b[keep[idx]]
      new_scores[:len(keep[idx]),3*k+1] = s[keep[idx]]
      #new tail class
      idx = [g for g in range(len(keep)) if keep[g] >=300]
      new_pred_boxes[:len(keep[idx]),12*k+8:12*k+12] = b[keep[idx]]
      new_scores[:len(keep[idx]),3*k+2] = s[keep[idx]]
      #new full length class = original
      new_pred_boxes[:,12*k+12:12*k+16] = pred_boxes[:,12*k+12:12*k+16]
      new_scores[:,3*k+3] = scores[:,3*k+3]
    if vis:
        im = cv2.imread(imdb.image_path_at(i))
        im2show = np.copy(im)

    for j in range(1, imdb.num_classes):
        inds = torch.nonzero(new_scores[:,j]>thresh).view(-1)
        # if there is det
        if inds.numel() > 0:
          cls_scores_1 = new_scores[:,j][inds]
          _, order = torch.sort(cls_scores_1, 0, True)
          if args.class_agnostic:
            cls_boxes_1 = new_pred_boxes[inds, :]
          else:
            cls_boxes_1 = new_pred_boxes[inds][:, j * 4:(j + 1) * 4]
          cls_dets_1 = torch.cat((cls_boxes_1, cls_scores_1.unsqueeze(1)), 1)
          cls_dets_1 = cls_dets_1[order]
          keep = nms(cls_boxes_1[order, :], cls_scores_1[order], cfg.TEST.NMS)
          cls_dets_1 = cls_dets_1[keep.view(-1).long()]

          if vis:
            im2show = vis_detections(im2show, imdb.classes[j], cls_dets.cpu().numpy(), 0.3)
          all_boxes_1[j][i] = cls_dets.cpu().numpy()
        else:
          all_boxes_1[j][i] = empty_array

    # Limit to max_per_image detections *over all classes*
    if max_per_image > 0:
        #print(all_boxes[3][i][:,-1])
        image_scores = np.hstack([all_boxes_1[j][i][:,-1]
                                  for j in range(1, imdb.num_classes)])
        if len(image_scores) > max_per_image:
            image_thresh = np.sort(image_scores)[-max_per_image]
            for j in range(1, imdb.num_classes):
                keep = np.where(all_boxes_1[j][i][:, -1] >= image_thresh)[0]
                all_boxes_1[j][i] = all_boxes_1[j][i][keep, :]

    misc_toc = time.time()
    nms_time = misc_toc - misc_tic

    sys.stdout.write('im_detect: {:d}/{:d} {:.3f}s {:.3f}s   \r' \
        .format(i + 1, len(dataloader1), detect_time, nms_time))
    sys.stdout.flush()

  torch.cuda.empty_cache()
  with open(det_file, 'wb') as f:
      pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)

  print('Evaluating detections')
  end = time.time()
  #print(len(all_boxes))
  #print(len(all_boxes_1[0]))
  for a in range(len(all_boxes)):
    all_boxes[a].extend(all_boxes_1[a])
    print(len(all_boxes[a]))
  print(new_indexes)
  #print(new_gt_boxes)
  imdb.evaluate_detections(all_boxes, output_dir, new_indexes, new_gt_boxes)

  print("test time: %0.4fs" % (end - start))


from __future__ import absolute_import  # be compatible with Python 2
from __future__ import division         # as there's no Python 3 on my lab's server
from __future__ import print_function

import align.detect_face
import facenet
import numpy as np
import tensorflow as tf
from scipy import misc
import os

data_dir = "../data"
image_index_file = os.path.join(data_dir, "new_image_index.txt")
image_dir = os.path.join(data_dir, "image_data")
face_feature_file = os.path.join(data_dir, "face_feature.npy")
face_info_file = os.path.join(data_dir, "face_info.txt")

facenet_model = "20180408-102900/"

def load_and_detect_faces(image_index_file, face_size, margin, gpu_memory_fraction=0.6):

    minsize = 20  # minimum size of face
    threshold = [0.6, 0.7, 0.7]  # three steps's threshold
    factor = 0.709  # scale factor
    info_list = []
    face_list = []
    face_num = 0
    image_num = 0
    with tf.Graph().as_default():
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_memory_fraction)
        sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, log_device_placement=False))
        with sess.as_default():
            pnet, rnet, onet = align.detect_face.create_mtcnn(sess, None)
    with open(image_index_file, mode="r") as image_index:
        for line in image_index:
            print("\rFaces Detected:{}/{}".format(face_num, image_num), end="")
            image_num += 1
            image_url, image_file, page_url, cv2_face_num = line.split()
            cv2_face_num = int(cv2_face_num)
            image_file = os.path.join(image_dir, image_file)
            image = misc.imread(image_file, mode="RGB")
            image_size = image.shape[0:2]
            bounding_boxes, _ = align.detect_face.detect_face(image, minsize, pnet, rnet, onet, threshold, factor)
            if len(bounding_boxes) < 1:
                print("\ncan't detect face, ignore ", image_file)
                continue
            for box in bounding_boxes:
                det = np.squeeze(box[0:4])
                bb = np.zeros(4, dtype=np.int32)
                bb[0] = np.maximum(det[0]-margin/2, 0)
                bb[1] = np.maximum(det[1]-margin/2, 0)
                bb[2] = np.minimum(det[2]+margin/2, image_size[1])
                bb[3] = np.minimum(det[3]+margin/2, image_size[0])
                cropped = image[bb[1]:bb[3], bb[0]:bb[2], :]
                aligned = misc.imresize(cropped, (face_size, face_size), interp='bilinear')
                prewhitened = facenet.prewhiten(aligned)
                face_list.append(prewhitened)
                face_num += 1
                info_list.append([image_url, image_file, page_url])
    face_array = np.stack(face_list)
    return face_array, info_list

def main():

    face_array, info_list = load_and_detect_faces(image_index_file, face_size=160, margin=44)
    face_embeddings = None
    face_num = len(face_array)
    assert face_num == len(info_list)
    with open(face_info_file, mode="w") as file:
        for info in info_list:
            file.write("{}\t{}\t{}\n".format(*info))
    face_embeddings = list()
    for i in range(0, face_num, 1000):  # avoid tf's OOM error
        slice_from, slice_to = i, min(i+1000, face_num)
        with tf.Graph().as_default():
            gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.6)
            with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, log_device_placement=False)) as sess:
                # Load the model
                facenet.load_model(facenet_model)
                # Get input and output tensors
                images_placeholder = tf.get_default_graph().get_tensor_by_name("input:0")
                embeddings = tf.get_default_graph().get_tensor_by_name("embeddings:0")
                phase_train_placeholder = tf.get_default_graph().get_tensor_by_name("phase_train:0")
                # Run forward pass to calculate embeddings
                feed_dict = {images_placeholder: face_array[slice_from:slice_to], phase_train_placeholder: False}
                print("Begins to compute embeddings...")
                face_embeddings.append(sess.run(embeddings, feed_dict=feed_dict))
                print(face_embeddings[-1].shape)
    face_embeddings = np.vstack(face_embeddings)
    assert face_num == len(face_embeddings), face_embeddings.shape
    np.save(face_feature_file, face_embeddings)

if __name__ == '__main__':

    main()