import os
import shutil
import numpy as np
from pyrevolve.custom_logging.logger import logger
import sys


class ExperimentManagement:
    # ids of robots in the name of all types of files are always phenotype ids, and the standard for id is 'robot_ID'

    def __init__(self, settings):
        self.settings = settings
        manager_folder = os.path.dirname(self.settings.manager)
        self._experiment_folder = os.path.join(manager_folder, 'data', self.settings.experiment_name, self.settings.run)
        self._data_folder = os.path.join(self._experiment_folder, 'data_fullevolution')
        self._gen_num = 0

    def create_exp_folders(self):
        if os.path.exists(self.experiment_folder):
            shutil.rmtree(self.experiment_folder)
        os.makedirs(self.experiment_folder)
        os.mkdir(self.data_folder)

        folders = ['genotypes', 'phenotypes', 'descriptors', 'objectives', 'fitness',
                   'battery', 'phenotype_images', 'failed_eval_robots']
        for folder in folders:
            os.mkdir(os.path.join(self.data_folder, folder))

    @property
    def experiment_folder(self):
        return self._experiment_folder

    @property
    def data_folder(self):
        return self._data_folder

    def export_genotype(self, individual):
        if self.settings.recovery_enabled:
            individual.export_genotype(self.data_folder)

    def export_phenotype(self, individual):
        if self.settings.export_phenotype:
            individual.export_phenotype(self.data_folder)

    def export_fitnesses(self, individuals):
        folder = self.data_folder
        for individual in individuals:
            individual.export_fitness(folder)

    def export_fitness(self, individual):
        folder = os.path.join(self.data_folder, 'fitness')
        individual.export_fitness(folder)

    def export_objectives(self, individual):
        folder = os.path.join(self.data_folder, 'objectives')
        individual.export_objectives(folder)

    def export_battery(self, individual):
        folder = os.path.join(self.data_folder, 'battery')
        individual.export_battery(folder)

    def export_behavior_measures(self, _id, measures):
        filename = os.path.join(self.data_folder, 'descriptors', f'behavior_desc_{_id}.txt')
        with open(filename, "w") as f:
            if measures is None:
                f.write(str(None))
            else:
                for key, val in measures.items():
                    f.write(f"{key} {val}\n")

    def export_phenotype_images(self, dirpath, individual):
        individual.phenotype.render_body(os.path.join(self.experiment_folder, dirpath, f'body_{individual.phenotype.id}.png'))
        individual.phenotype.render_brain(os.path.join(self.experiment_folder, dirpath, f'brain_{individual.phenotype.id}.png'))

    def export_failed_eval_robot(self, individual):
        individual.genotype.export_genotype(os.path.join(self.data_folder, 'failed_eval_robots', f'genotype_{individual.phenotype.id}.txt'))
        individual.phenotype.save_file(os.path.join(self.data_folder, 'failed_eval_robots', f'phenotype_{individual.phenotype.id}.yaml'))
        individual.phenotype.save_file(os.path.join(self.data_folder, 'failed_eval_robots', f'phenotype_{individual.phenotype.id}.sdf'), conf_type='sdf')

    def export_snapshots(self, individuals, gen_num):
        self._gen_num = gen_num
        if self.settings.recovery_enabled:
            path = os.path.join(self.experiment_folder, f'selectedpop_{gen_num}')
            if os.path.exists(path):
                shutil.rmtree(path)
            os.mkdir(path)
            for ind in individuals:
                self.export_phenotype_images(f'selectedpop_{str(gen_num)}', ind)
            logger.info(f'Exported snapshot {str(gen_num)} with {str(len(individuals))} individuals')

    def experiment_is_new(self):
        if not os.path.exists(self.experiment_folder):
            return True
        path, dirs, files = next(os.walk(os.path.join(self.data_folder, 'fitness')))
        if len(files) == 0:
            return True
        else:
            return False

    def read_recovery_state(self, population_size, offspring_size):
        snapshots = []

        for r, d, f in os.walk(self.experiment_folder):
            for dir in d:
                if 'selectedpop' in dir:
                    exported_files = len([name for name in os.listdir(os.path.join(self.experiment_folder, dir)) if os.path.isfile(os.path.join(self.experiment_folder, dir, name))])
                    if exported_files == (population_size * 2): # body and brain files
                        snapshots.append(int(dir.split('_')[1]))

        if len(snapshots) > 0:
            # the latest complete snapshot
            last_snapshot = np.sort(snapshots)[-1]
            # number of robots expected until the snapshot
            n_robots = population_size + last_snapshot * offspring_size
        else:
            last_snapshot = -1
            n_robots = 0

        robot_ids = []

        for r, d, f in os.walk(os.path.join(self.data_folder, 'fitness')):
            for file in f:
                robot_ids.append(int(file.split('.')[0].split('_')[-1]))
        last_id = np.sort(robot_ids)[-1]

        # if there are more robots to recover than the number expected in this snapshot
        if last_id > n_robots:
            # then recover also this partial offspring
            has_offspring = True
        else:
            has_offspring = False

        return last_snapshot, has_offspring, last_id+1

    def plot_path(self, data_source: str, filename: str, file_extension=".png"):
        data_folder = os.path.join(self._data_folder, data_source)
        if not os.path.exists(data_folder):
            os.mkdir(data_folder)

        return os.path.join(data_folder, filename + str(self._gen_num) + file_extension)

[(13, 'os.path.dirname', 'os.path.dirname', False, 'import os\n'), (14, 'os.path.join', 'os.path.join', False, 'import os\n'), (15, 'os.path.join', 'os.path.join', False, 'import os\n'), (19, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (21, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (22, 'os.mkdir', 'os.mkdir', False, 'import os\n'), (51, 'os.path.join', 'os.path.join', False, 'import os\n'), (55, 'os.path.join', 'os.path.join', False, 'import os\n'), (59, 'os.path.join', 'os.path.join', False, 'import os\n'), (63, 'os.path.join', 'os.path.join', False, 'import os\n'), (103, 'os.walk', 'os.walk', False, 'import os\n'), (136, 'os.path.join', 'os.path.join', False, 'import os\n'), (20, 'shutil.rmtree', 'shutil.rmtree', False, 'import shutil\n'), (72, 'os.path.join', 'os.path.join', False, 'import os\n'), (73, 'os.path.join', 'os.path.join', False, 'import os\n'), (76, 'os.path.join', 'os.path.join', False, 'import os\n'), (77, 'os.path.join', 'os.path.join', False, 'import os\n'), (78, 'os.path.join', 'os.path.join', False, 'import os\n'), (83, 'os.path.join', 'os.path.join', False, 'import os\n'), (84, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (86, 'os.mkdir', 'os.mkdir', False, 'import os\n'), (92, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (121, 'os.path.join', 'os.path.join', False, 'import os\n'), (124, 'numpy.sort', 'np.sort', True, 'import numpy as np\n'), (137, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (138, 'os.mkdir', 'os.mkdir', False, 'import os\n'), (27, 'os.path.join', 'os.path.join', False, 'import os\n'), (85, 'shutil.rmtree', 'shutil.rmtree', False, 'import shutil\n'), (94, 'os.path.join', 'os.path.join', False, 'import os\n'), (112, 'numpy.sort', 'np.sort', True, 'import numpy as np\n'), (106, 'os.path.join', 'os.path.join', False, 'import os\n'), (106, 'os.path.join', 'os.path.join', False, 'import os\n')]