# import matplotlib
# matplotlib.use('MacOSX')

from typing import List, Optional

from habitat.config import Config
from habitat.datasets.pointnav.pointnav_dataset import PointNavDatasetV1
from habitat.core.registry import registry

import argparse
import habitat
import random
import numpy
import os
import time
import cv2
from habitat.tasks.nav.shortest_path_follower import ShortestPathFollower
from habitat.utils.visualizations import maps
from myhabitatagent import DSLAMAgent
import numpy as np
import tensorflow as tf

from arguments import parse_args

import habitat_utils

# This is a hack attempt to address error "BrokenPipeError: [Errno 32] Broken pipe". My suspicion is that the
# planner process can be alive (and write to stdout) even after all episodes are done. The output is possibly
# redirected in the eval script, which creates a problem similar to this:
# https://stackoverflow.com/questions/11423225/why-does-my-python3-script-balk-at-piping-its-output-to-head-or-tail-sys-module
# The issue was discovered with ./test_locally_pointnav_rgbs.sh with an TIMEOUT of 1sec, where agent gives up all episdoes.
# import signal
# signal.signal(signal.SIGPIPE, signal.SIG_DFL)
# signal.signal(signal.SIG, signal.SIG_DFL)



@registry.register_dataset(name="PointNavCustomSplit-v3")
class PointNavCustomSplitDatasetV3(PointNavDatasetV1):
    def __init__(self, config: Optional[Config] = None) -> None:
        super(PointNavCustomSplitDatasetV3, self).__init__(config)

        # Filter for custom split
        if config.CUSTOM_SPLIT:
            HABITAT_VAL3_SCENES = ('Cantwell', 'Denmark', 'Eastville',
                                   'Edgemere', 'Elmira', 'Eudora',
                                   'Greigsville')

            if config.CUSTOM_SPLIT == 'habitat-test3':
                is_valid = lambda scene_id: (scene_id not in HABITAT_VAL3_SCENES)
            elif config.CUSTOM_SPLIT == 'habitat-val3':
                is_valid = lambda scene_id: (scene_id in HABITAT_VAL3_SCENES)
            else:
                raise ValueError("Unknown custom split %s"%config.CUSTOM_SPLIT)

            new_episodes = []
            for ep in self.episodes:
                scene_id = habitat_utils.get_model_id_from_episode(ep)
                if is_valid(scene_id):
                    print (scene_id)
                    new_episodes.append(ep)
            print("Total %d --> %d" % (len(self.episodes), len(new_episodes)))
            self.episodes = new_episodes




class RandomAgent(habitat.Agent):
    def __init__(self, task_config: habitat.Config, params=None):
        self._POSSIBLE_ACTIONS = task_config.TASK.POSSIBLE_ACTIONS
        print (params)

    def reset(self):
        print ("reseting..")
        self.step_i = 0
        pass

    def act(self, observations):
        self.step_i += 1
        if self.step_i > 100:
            action = 0
        else:
            action = numpy.random.choice([1, 2, 3])
        print (self.step_i, action)
        return {"action": action}


class ShortestPathAgent(habitat.Agent):
    def __init__(self, task_config: habitat.Config, env=None):
        self._POSSIBLE_ACTIONS = task_config.TASK.POSSIBLE_ACTIONS
        self.step_i = 0
        self.episode_i = -2
        self.env = env
        self.follower = ShortestPathFollower(env._sim, 0.36/2., False)
        self.reset()

    def reset(self):
        self.step_i = 0
        self.episode_i += 1
        print ("Resetting agent %d. Scene %s."%(self.episode_i, self.env._sim._current_scene))

    def act(self, observations):
        # print (observations.keys())
        # import ipdb; ipdb.set_trace()

        goal_pos = self.env.current_episode.goals[0].position
        best_action = self.follower.get_next_action(goal_pos)

        if self.episode_i == 0:
            cv2.imwrite('./temp/ep%d-step%d.png'%(self.episode_i, self.step_i), observations['rgb'])
        if self.step_i == 0:
            top_down_map = maps.get_topdown_map(
                self.env.sim, map_resolution=(5000, 5000)
            )
            import matplotlib.pyplot as plt
            plt.imshow(top_down_map)
            plt.show()
            plt.waitforbuttonpress(0.001)

        #     global_map = observations['top_down_map']
        #     cv2.imwrite('./temp/map-%d.png'%self.episode_i, global_map)
        self.step_i += 1

        return {"action": best_action}   # 0: stop, forward, left, right
        # return {"action": numpy.random.choice(self._POSSIBLE_ACTIONS)}


def main():
    params = parse_args(default_files=('./habitat_submission.conf', ))
    is_submission = (params.habitat_eval not in ['localtest', 'gendata'])

    if params.seed > 0:
        np.random.seed(params.seed)
        random.seed(params.seed)

    # Update for 2021 change in environment variable naming
    if os.environ["TRACK_CONFIG_FILE"]:
        print ("Overwriting CHALLENGE_CONFIG_FILE with %s"%os.environ["TRACK_CONFIG_FILE"])
        os.environ["CHALLENGE_CONFIG_FILE"] = os.environ["TRACK_CONFIG_FILE"]

    if not is_submission:
        if params.habitat_config_file != "":
            os.environ["CHALLENGE_CONFIG_FILE"] = params.habitat_config_file
        elif params.overwrite_agent_noise >= 0.:
            assert  params.overwrite_agent_noise == 0.  # need a separate file otherwise
            print("Overwriting agent noise %f" % params.overwrite_agent_noise)
            os.environ["CHALLENGE_CONFIG_FILE"] = './configs/challenge_pointnav_supervised_nonoise.yaml'
        else:
            os.environ["CHALLENGE_CONFIG_FILE"] = './configs/challenge_pointnav_supervised.yaml'

    # parser = argparse.ArgumentParser()
    # parser.add_argument("--evaluation", type=str, required=True, choices=["local", "remote"])
    # parser.add_argument("--config", type=str, default='', required=False)
    # args = parser.parse_args()
    #
    # if args.config != '':
    #     os.environ["CHALLENGE_CONFIG_FILE"] = args.config

    config_paths = os.environ["CHALLENGE_CONFIG_FILE"]
    config = habitat.get_config(config_paths)

    if config.SIMULATOR.TYPE.startswith('Spot'):
        # This will register the custom sim class
        from spot import spotsim

        # Update compressed episode json
        jsonfile = config.DATASET.DATA_PATH
        assert jsonfile.endswith('.gz')
        jsonfile = jsonfile[:-3]
        import subprocess
        subprocess.call(['gzip -kf ' + jsonfile], shell=True)
        # import ipdb as pdb
        # pdb.set_trace()

        # from habitat.sims.pyrobot import pyrobot

    print ("Using config file(s): %s"%(str(config_paths)))
    logdir = params.logpath if params.logpath != '' \
        else './temp/evals/{}-{}'.format(time.strftime('%m-%d-%H-%M-%S', time.localtime()), params.name)  # this is only used for videos now, should pass it in to challenge eval
    os.makedirs(logdir, exist_ok=True)
    # agent = RandomAgent(task_config=config)

    if params.habitat_eval in ["localtest", "gendata"]:
        # # TODO Try to add additional noiseless sensors. Changes here have no effect!! Modify config file instead.
        # config.defrost()
        # # config.SIMULATOR.AGENT_0.SENSORS.append("NOISELESS_RGB_SENSOR")
        # # config.SIMULATOR.NOISELESS_RGB_SENSOR = habitat.Config(config.SIMULATOR.RGB_SENSOR)
        # # config.SIMULATOR.NOISELESS_RGB_SENSOR.NOISE_MODEL = ""
        # # config.SIMULATOR.NOISELESS_RGB_SENSOR.NOISE_MODEL_KWARGS = habitat.Config()
        # # config.SIMULATOR.AGENT_0.SENSORS.append("NOISELESS_DEPTH_SENSOR")
        # # config.SIMULATOR.NOISELESS_DEPTH_SENSOR = habitat.Config(config.SIMULATOR.DEPTH_SENSOR)
        # # config.SIMULATOR.NOISELESS_DEPTH_SENSOR.NOISE_MODEL = ""
        # # config.SIMULATOR.NOISELESS_DEPTH_SENSOR.NOISE_MODEL_KWARGS = habitat.Config()
        # config.SIMULATOR.AGENT_0.SENSORS.append("NOISELESS_RGB_SENSOR")
        # config.SIMULATOR.NOISELESS_RGB_SENSOR = habitat.Config(config.SIMULATOR.RGB_SENSOR)
        # config.SIMULATOR.NOISELESS_RGB_SENSOR.TYPE = "SecondHabitatSimRGBSensor"
        # config.SIMULATOR.NOISELESS_RGB_SENSOR.NOISE_MODEL = ""
        # config.SIMULATOR.NOISELESS_RGB_SENSOR.NOISE_MODEL_KWARGS = habitat.Config()
        # config.freeze()

        challenge = habitat.Challenge(eval_remote=False)
        env = challenge._env

        grid_cell_size = 0.05  # 5cm
        map_size = (maps.COORDINATE_MAX - maps.COORDINATE_MIN) / grid_cell_size
        assert config.TASK.TOP_DOWN_MAP.MAP_RESOLUTION == int(map_size)

        if params.seed > 0:
            env._sim.seed(params.seed)
            env.seed(params.seed)

        # agent = ShortestPathAgent(task_config=config, env=env)

        if params.habitat_eval == "gendata":
            assert len(params.data_map_sizes) > 0

            # Initialize writers
            data_filenames = [os.path.join(
                logdir, "agentplandata.m%d.%d-%d.tfrecords"%(map_size, params.skip_first_n, params.num_episodes))
                for map_size in params.data_map_sizes]

            # with statement for variable number of items
            from contextlib import ExitStack
            with ExitStack() as stack:
                tfwriters = [stack.enter_context(tf.python_io.TFRecordWriter(data_filename)) for data_filename in data_filenames]
                agent = DSLAMAgent(task_config=config, params=params, env=env, logdir=logdir, tfwriters=tfwriters)
                challenge.submit(agent, num_episodes=params.num_episodes, skip_first_n=params.skip_first_n)
        else:
            agent = DSLAMAgent(task_config=config, params=params, env=env, logdir=logdir)
            challenge.submit(agent, num_episodes=params.num_episodes, skip_first_n=params.skip_first_n)

    elif params.habitat_eval == "local":
        challenge = habitat.Challenge(eval_remote=False)

        # if params.seed > 0:
        #     challenge._env._sim.seed(params.seed)
        #     challenge._env.seed(params.seed)

        agent = DSLAMAgent(task_config=config, params=params, env=None, logdir=logdir)
        # agent = RandomAgent(task_config=config, params=params)
        challenge.submit(agent)  # , num_episodes=params.num_episodes)

    else:
        challenge = habitat.Challenge(eval_remote=True)
        agent = DSLAMAgent(task_config=config, params=params, env=None, logdir=logdir)
        # agent = RandomAgent(task_config=config, params=params)
        challenge.submit(agent)


# TODO we can add a second RGB sensor this way, but i have not implemented the logic to save to data and to load such data.
# from habitat.sims.habitat_simulator import habitat_simulator
# from typing import Any, List, Optional, Union
#
# # Define the sensor and register it with habitat
# # For the sensor, we will register it with a custom name
# @habitat.registry.register_sensor(name="SecondHabitatSimRGBSensor")
# class SecondHabitatSimRGBSensor(habitat_simulator.HabitatSimRGBSensor):
#     # Defines the name of the sensor in the sensor suite dictionary
#     def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
#         return "rgb2"
#
#
# @habitat.registry.register_sensor(name="SecondHabitatSimDepthSensor")
# class SecondHabitatSimDepthSensor(habitat_simulator.HabitatSimDepthSensor):
#     # Defines the name of the sensor in the sensor suite dictionary
#     def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
#         return "depth2"


if __name__ == "__main__":
    main()
