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, 25, 13, 63), 'os.path.dirname', 'os.path.dirname', ({(13, 41, 13, 62): 'self.settings.manager'}, {}), '(self.settings.manager)', False, 'import os\n'), ((14, 34, 14, 120), 'os.path.join', 'os.path.join', ({(14, 47, 14, 61): 'manager_folder', (14, 63, 14, 69): '"""data"""', (14, 71, 14, 100): 'self.settings.experiment_name', (14, 102, 14, 119): 'self.settings.run'}, {}), "(manager_folder, 'data', self.settings.experiment_name, self.\n    settings.run)", False, 'import os\n'), ((15, 28, 15, 87), 'os.path.join', 'os.path.join', ({(15, 41, 15, 64): 'self._experiment_folder', (15, 66, 15, 86): '"""data_fullevolution"""'}, {}), "(self._experiment_folder, 'data_fullevolution')", False, 'import os\n'), ((19, 11, 19, 49), 'os.path.exists', 'os.path.exists', ({(19, 26, 19, 48): 'self.experiment_folder'}, {}), '(self.experiment_folder)', False, 'import os\n'), ((21, 8, 21, 43), 'os.makedirs', 'os.makedirs', ({(21, 20, 21, 42): 'self.experiment_folder'}, {}), '(self.experiment_folder)', False, 'import os\n'), ((22, 8, 22, 34), 'os.mkdir', 'os.mkdir', ({(22, 17, 22, 33): 'self.data_folder'}, {}), '(self.data_folder)', False, 'import os\n'), ((51, 17, 51, 58), 'os.path.join', 'os.path.join', ({(51, 30, 51, 46): 'self.data_folder', (51, 48, 51, 57): '"""fitness"""'}, {}), "(self.data_folder, 'fitness')", False, 'import os\n'), ((55, 17, 55, 61), 'os.path.join', 'os.path.join', ({(55, 30, 55, 46): 'self.data_folder', (55, 48, 55, 60): '"""objectives"""'}, {}), "(self.data_folder, 'objectives')", False, 'import os\n'), ((59, 17, 59, 58), 'os.path.join', 'os.path.join', ({(59, 30, 59, 46): 'self.data_folder', (59, 48, 59, 57): '"""battery"""'}, {}), "(self.data_folder, 'battery')", False, 'import os\n'), ((63, 19, 63, 92), 'os.path.join', 'os.path.join', ({(63, 32, 63, 48): 'self.data_folder', (63, 50, 63, 63): '"""descriptors"""', (63, 65, 63, 91): 'f"""behavior_desc_{_id}.txt"""'}, {}), "(self.data_folder, 'descriptors', f'behavior_desc_{_id}.txt')", False, 'import os\n'), ((103, 23, 103, 54), 'os.walk', 'os.walk', ({(103, 31, 103, 53): 'self.experiment_folder'}, {}), '(self.experiment_folder)', False, 'import os\n'), ((136, 22, 136, 66), 'os.path.join', 'os.path.join', ({(136, 35, 136, 52): 'self._data_folder', (136, 54, 136, 65): 'data_source'}, {}), '(self._data_folder, data_source)', False, 'import os\n'), ((20, 12, 20, 49), 'shutil.rmtree', 'shutil.rmtree', ({(20, 26, 20, 48): 'self.experiment_folder'}, {}), '(self.experiment_folder)', False, 'import shutil\n'), ((72, 41, 72, 125), 'os.path.join', 'os.path.join', ({(72, 54, 72, 76): 'self.experiment_folder', (72, 78, 72, 85): 'dirpath', (72, 87, 72, 124): 'f"""body_{individual.phenotype.id}.png"""'}, {}), "(self.experiment_folder, dirpath,\n    f'body_{individual.phenotype.id}.png')", False, 'import os\n'), ((73, 42, 73, 127), 'os.path.join', 'os.path.join', ({(73, 55, 73, 77): 'self.experiment_folder', (73, 79, 73, 86): 'dirpath', (73, 88, 73, 126): 'f"""brain_{individual.phenotype.id}.png"""'}, {}), "(self.experiment_folder, dirpath,\n    f'brain_{individual.phenotype.id}.png')", False, 'import os\n'), ((76, 44, 76, 139), 'os.path.join', 'os.path.join', ({(76, 57, 76, 73): 'self.data_folder', (76, 75, 76, 95): '"""failed_eval_robots"""', (76, 97, 76, 138): 'f"""genotype_{individual.phenotype.id}.txt"""'}, {}), "(self.data_folder, 'failed_eval_robots',\n    f'genotype_{individual.phenotype.id}.txt')", False, 'import os\n'), ((77, 39, 77, 136), 'os.path.join', 'os.path.join', ({(77, 52, 77, 68): 'self.data_folder', (77, 70, 77, 90): '"""failed_eval_robots"""', (77, 92, 77, 135): 'f"""phenotype_{individual.phenotype.id}.yaml"""'}, {}), "(self.data_folder, 'failed_eval_robots',\n    f'phenotype_{individual.phenotype.id}.yaml')", False, 'import os\n'), ((78, 39, 78, 135), 'os.path.join', 'os.path.join', ({(78, 52, 78, 68): 'self.data_folder', (78, 70, 78, 90): '"""failed_eval_robots"""', (78, 92, 78, 134): 'f"""phenotype_{individual.phenotype.id}.sdf"""'}, {}), "(self.data_folder, 'failed_eval_robots',\n    f'phenotype_{individual.phenotype.id}.sdf')", False, 'import os\n'), ((83, 19, 83, 81), 'os.path.join', 'os.path.join', ({(83, 32, 83, 54): 'self.experiment_folder', (83, 56, 83, 80): 'f"""selectedpop_{gen_num}"""'}, {}), "(self.experiment_folder, f'selectedpop_{gen_num}')", False, 'import os\n'), ((84, 15, 84, 35), 'os.path.exists', 'os.path.exists', ({(84, 30, 84, 34): 'path'}, {}), '(path)', False, 'import os\n'), ((86, 12, 86, 26), 'os.mkdir', 'os.mkdir', ({(86, 21, 86, 25): 'path'}, {}), '(path)', False, 'import os\n'), ((92, 15, 92, 53), 'os.path.exists', 'os.path.exists', ({(92, 30, 92, 52): 'self.experiment_folder'}, {}), '(self.experiment_folder)', False, 'import os\n'), ((121, 31, 121, 72), 'os.path.join', 'os.path.join', ({(121, 44, 121, 60): 'self.data_folder', (121, 62, 121, 71): '"""fitness"""'}, {}), "(self.data_folder, 'fitness')", False, 'import os\n'), ((124, 18, 124, 36), 'numpy.sort', 'np.sort', ({(124, 26, 124, 35): 'robot_ids'}, {}), '(robot_ids)', True, 'import numpy as np\n'), ((137, 15, 137, 42), 'os.path.exists', 'os.path.exists', ({(137, 30, 137, 41): 'data_folder'}, {}), '(data_folder)', False, 'import os\n'), ((138, 12, 138, 33), 'os.mkdir', 'os.mkdir', ({(138, 21, 138, 32): 'data_folder'}, {}), '(data_folder)', False, 'import os\n'), ((27, 21, 27, 59), 'os.path.join', 'os.path.join', ({(27, 34, 27, 50): 'self.data_folder', (27, 52, 27, 58): 'folder'}, {}), '(self.data_folder, folder)', False, 'import os\n'), ((85, 16, 85, 35), 'shutil.rmtree', 'shutil.rmtree', ({(85, 30, 85, 34): 'path'}, {}), '(path)', False, 'import shutil\n'), ((94, 41, 94, 82), 'os.path.join', 'os.path.join', ({(94, 54, 94, 70): 'self.data_folder', (94, 72, 94, 81): '"""fitness"""'}, {}), "(self.data_folder, 'fitness')", False, 'import os\n'), ((112, 28, 112, 46), 'numpy.sort', 'np.sort', ({(112, 36, 112, 45): 'snapshots'}, {}), '(snapshots)', True, 'import numpy as np\n'), ((106, 70, 106, 111), 'os.path.join', 'os.path.join', ({(106, 83, 106, 105): 'self.experiment_folder', (106, 107, 106, 110): 'dir'}, {}), '(self.experiment_folder, dir)', False, 'import os\n'), ((106, 131, 106, 178), 'os.path.join', 'os.path.join', ({(106, 144, 106, 166): 'self.experiment_folder', (106, 168, 106, 171): 'dir', (106, 173, 106, 177): 'name'}, {}), '(self.experiment_folder, dir, name)', False, 'import os\n')]