research.py

import pickle
from datetime import date
from typing import Sequence, Mapping, TypedDict

from comsdk.comaux import *
from comsdk.communication import BaseCommunication, LocalCommunication, SshCommunication, Host
from comsdk.distributed_storage import *
from comsdk.edge import Func, Edge, dummy_predicate
from comsdk.graph import Graph, State

CopiesList = TypedDict('CopiesList', {'path': str, 'new_name': str})


class Research:

    """
    Class Research is a representation of a group of different calculations collected into what we call a Research.
    Each ''calculation'' corresponds to the launch of a graph-based scenario which produces a set of files which we
    treat as the results of the calculation. We thus call such a calculation a task. Therefore, a Research is a
    collection of tasks. Each task is associated with a single directory (within the code, it may be denoted as
    task_dir, if only directory name is of interest, or task_path, if the absolute path is of interest) whose name has
    a very simple structure, @number@-@long_name@, so that each task is associated with its own unique number (also
    called task_number within the code). Normally, one should use the task number to get any task-related information.
    All the tasks are located in the research directory whose the local (remote) absolute path is set by the class
    property local_research_path (remote_research_path). The research directory has the following pattern:
    @date@_@long_name@. Finally, we associate a short Research ID with each Research. The described structure is
    independent of where these directories are located. It is assumed that there is a local root for research and
    its remote analog. The latter should be available via any protocol supported by communication module. Class Research
    can thus be set up in two regimes: local (remote_comm is None) and local-remote (remote_comm is not None).

    Typically, one should construct an instance of Research based on the configuration file called config_research.json.
    There are two static functions for this purpose: Research.open() and Research.create(). The former creates an
    instance of Research based on the existing Research (one should pass its Research ID to open()) described in the
    configuration file and the latter creates a new Research (thus, making a new directory in the local filesystem) and
    adds all the necessary information about it in the configuration file. Also, any Research instance is automatically
    augmented by the properties listed in 'RESEARCH_PROPS' dictionary in the configuration file.

    For the Research constructor to understand where all the research directories are located, one must supply (either
    directly in the constructor or in the configuration file) the potential root paths for the search (both for the
    local and remote machines if the latter is specified). The first path in the list of the potential root paths is
    called the default root path. A new Research will be created in the default path.

    Note that different tasks belonging to the same research (i.e., they are associated with the same Research ID) may
    be located at different root paths. When creating a new task, it will be located in the default root path.

    .. todo::
        Some way for saving auxiliary information about research and tasks (task date and description, for example)
        should be implemented. Possibly, the same should be done for launcher scripts.

    """
    def __init__(self, name: str,
                 continuing=False,
                 local_research_roots: Optional[Sequence[str]] = None,
                 remote_comm: Optional[BaseCommunication] = None,
                 remote_research_root: Optional[str] = None):
        """
        :param name: research description (if continuing == False) or research directory (if continuing == True)
        :param continuing: if False, the Research with be read from the root path. Otherwise, a new one will be created
        :param local_research_roots: a list of local paths where research directories are searched for
        :param remote_comm: BaseCommunication instance used for communication with remote machine
        :param remote_research_root: path on the remote machine where research directories are searched for
        """
        self._local_research_root = local_research_roots[0]
        self._local_root = os.path.dirname(self._local_research_root)
        self._remote_research_root = remote_research_root
        self._tasks_number = 0
        self._local_comm = LocalCommunication(Host())  # local communication created automatically, no need to pass it
        self._remote_comm = remote_comm
        self._distr_storage = DistributedStorage(local_research_roots, prior_storage_index=0)
        self._local_research_path = None
        if not continuing:
            # interpret name as name without date
            self._research_dir = make_suitable_research_dir(name)
            if self._distr_storage.get_dir_path(self._research_dir) is not None:
                raise ResearchAlreadyExists("Research with name '{}' already exists, "
                                            "choose another name".format(self._research_dir))
            self._local_research_path = self._distr_storage.make_dir(self._research_dir)
            print('Started new research at {}'.format(self._local_research_path))
        else:
            # interpret name as the full research id
            self._research_dir = name
            self._local_research_path = self._load_research_data()

    @classmethod
    def open(cls, research_id: str,
             remote_comm: Optional[BaseCommunication] = None):
        """
        :param research_id: Research ID used to find a relevant research
        :param remote_comm: BaseCommunication instance used for communication with remote machine
        :return: new Research instance
        """
        with open('config_research.json', 'r') as f:
            conf = json.load(f)
        res = Research(conf['RESEARCH'][research_id],
                       continuing=True,
                       local_research_roots=conf['LOCAL_HOST']['research_roots'],
                       remote_comm=remote_comm,
                       remote_research_root=conf['REMOTE_HOSTS'][remote_comm.machine_name]['research_root']
                                            if remote_comm is not None else None)
        res._add_properties(conf['RESEARCH_PROPS'])
        return res

    @classmethod
    def create(cls, new_research_id: str, new_research_descr: str,
               remote_comm: Optional[BaseCommunication] = None):
        """
        :param new_research_id: Research ID (short name for this research)
        :param new_research_descr: relatively long research name
        :param remote_comm: BaseCommunication instance used for communication with remote machine
        :return: new Research instance
        """
        with open('config_research.json', 'r+') as f:
            conf = json.load(f)
            conf['RESEARCH'][new_research_id] = make_suitable_research_dir(new_research_descr)
            f.seek(0)
            json.dump(conf, f, indent=4)
            f.truncate()
        res = Research(new_research_descr,
                       continuing=False,
                       local_research_roots=conf['LOCAL_HOST']['research_roots'],
                       remote_comm=remote_comm,
                       remote_research_root=conf['REMOTE_HOSTS'][remote_comm.machine_name]['research_root']
                                            if remote_comm is not None else None)
        res._add_properties(conf['RESEARCH_PROPS'])
        return res

    @property
    def local_research_path(self) -> str:
        return self._local_research_path

    @property
    def remote_research_path(self) -> str:
        return os.path.join(self._remote_research_root, self._research_dir)

    @property
    def local_root(self) -> str:
        return self._local_root

    @property
    def research_dir(self) -> str:
        return self._research_dir

    def __getstate__(self) -> dict:
        return {
            'research_dir': self._research_dir,
            'local_research_path': self._local_research_root,
            'remote_research_path': self._remote_research_root,
            'remote_comm': self._remote_comm.__getstate__(),
        }

    def __setstate__(self, state):
        self._tasks_number = 0
        self._local_comm = LocalCommunication(Host())
        self._local_research_root = state['local_research_path']
        self._remote_research_root = state['remote_research_path']
        self._remote_comm = None
        if state['remote_comm'] is not None:
            self._remote_comm = SshCommunication.__new__(SshCommunication)
            self._remote_comm.__setstate__(state['remote_comm'])
        self._distr_storage = DistributedStorage((self._local_research_root,), prior_storage_index=0)
        self._research_dir = state['research_dir']
        self._research_path = self._load_research_data()

    def _add_properties(self, props: Mapping[str, Any]) -> None:
        for prop_name, prop_value in props.items():
            self.__setattr__(prop_name, prop_value)

    def _load_research_data(self) -> str:
        # find corresponding date/name
        # construct object from all data inside
        research_path = self._distr_storage.get_dir_path(self._research_dir)
        if research_path is None:
            raise ResearchDoesNotExist("Research '{}' does not exist".format(self._research_dir))
#        if research_path is None:
#            # assume date was omitted in research id
#            regexp_for_search = '^(?P<year>\d+)-(?P<month>\d+)-(?P<day>\d+)_{}'.format(self._research_name)
#            research_path, dir_params = self._distr_storage.find_dir_by_named_regexp('', regexp_for_search)
#            if dir_params is None:
#                raise ResearchDoesNotExist("Research '{}' does not exist".format(self._research_name))
#            self._research_name = '{}-{}-{}_{}'.format(dir_params['year'], dir_params['month'], dir_params['day'],
#                                                       self._research_name)

        print('Loaded research at {}'.format(research_path))

        # determine maximum task number to set the number for the next possible task
        dirnames, _ = self._distr_storage.listdir(self._research_dir)
        self._tasks_number = 0
        for dir_ in dirnames:
            if dir_ != 'report':
                task_number, _ = split_task_dir(dir_)
                if task_number > self._tasks_number:
                    self._tasks_number = task_number
        self._tasks_number += 1
        print('Number of tasks in the current research: {}'.format(self._tasks_number))
        return research_path

    def create_task(self, name: str) -> int:
        """
        Creates a new task in the current research making a new local directory

        :param name: task name
        :return: task number
        """
        task_number = self._get_next_task_number()
        local_task_dir = self._make_task_path(task_number, name)
        os.mkdir(local_task_dir)
        return task_number

    def grab_task_results(self, task_number: int,
                          copies_list: Optional[Sequence[CopiesList]] = None):
        """
        Moves task content from the remote machine to the local one. Locally, the task content will appear in the task
        directory located in the research directory.

        :param task_number: task number
        :param copies_list: a list defining which objects we wish to copy from the remote machine. It consists of
        dictionaries each having keys 'path' (path of object we wish to copy relative to the task directory) and
        'new_name' (path of this object on the local machine relative to the task directory)
        """
        task_results_local_path = self.get_task_path(task_number)
        task_results_remote_path = self.get_task_path(task_number, self._remote_comm.host)
        if copies_list is None: # copy all data
            paths = self._remote_comm.listdir(task_results_remote_path)
            for file_or_dir in paths:
                self._remote_comm.copy('/'.join((task_results_remote_path, file_or_dir)), task_results_local_path,
                                       'from_remote', show_msg=True)
        else:
            for copy_target in copies_list:
                # we consider copy targets as relative to task's dir
                remote_copy_target_path = '/'.join((task_results_remote_path, copy_target['path']))
                self._remote_comm.copy(remote_copy_target_path, task_results_local_path, 'from_remote', show_msg=True)
                if 'new_name' in copy_target:
                    os.rename(os.path.join(task_results_local_path, os.path.basename(copy_target['path'])),
                              os.path.join(task_results_local_path, copy_target['new_name']))

    def _make_task_path(self, task_number: int, task_name: str, at_remote_host=False) -> str:
        task_path = None
        task_dir = get_task_full_name(task_number, task_name)
        if at_remote_host:
            task_path = os.path.join(self._remote_research_root, self._research_dir, task_dir)
        else:
            task_path = os.path.join(self._local_research_path, task_dir)
        return task_path

    def get_task_path(self, task_number: int, at_remote_host=False) -> str:
        """
        Return absolute task path based on its number

        :param task_number: task number
        :param at_remote_host: return the path on the remote machine (if True) or on the local one (if False)
        :return: absolute task path
        """
        task_path = None
        task_name = self._get_task_name_by_number(task_number)
        rel_task_dir = os.path.join(self._research_dir, get_task_full_name(task_number, task_name))
        if at_remote_host:
            if self._remote_comm is None:
                raise ValueError('Cannot get a task path on the remote: remote communication is not set up')
            task_path = '{}/{}'.format(self._remote_research_root, rel_task_dir)
        else:
            task_path = self._distr_storage.get_dir_path(rel_task_dir)
        return task_path

    def dump_object(self, task_number: int, obj: object, obj_name: str) -> None:
        """
        Dumps any python object (using pickle) to the binary file, named obj_name + '.pyo', in the task directory
        associated with the task number

        :param task_number: task number
        :param obj: any python object
        :param obj_name: file name to which obj will be saved (without extension)
        """

        print('Dumping ' + obj_name)
        f = open(os.path.join(self.get_task_path(task_number), obj_name + '.pyo'), 'w')
        pickle.dump(obj, f)
        f.close()

    def load_object(self, task_number: int, obj_name: str):
        """
        Load any python object dumped using pickle from the binary file, named obj_name + '.pyo' and located in the task
        directory associated with the task number

        :param task_number: task number
        :param obj_name: file name from which obj will be loaded (without extension)
        :return: python object
        """
        print('Loading ' + obj_name)
        f = open(os.path.join(self.get_task_path(task_number), obj_name + '.pyo'), 'r')
        obj = pickle.load(f)
        f.close()
        return obj

    def _get_next_task_number(self) -> int:
        self._tasks_number += 1
        return self._tasks_number - 1

    def _get_task_name_by_number(self, task_number: int) -> str:
        find_data = self._distr_storage.find_dir_by_named_regexp(self._research_dir,
                                                                 '^{}-(?P<task_name>\S+)'.format(task_number))
        if find_data is None:
            raise Exception("No task with number '{}' is found".format(task_number))
        return find_data[1]['task_name']


class ResearchAlreadyExists(Exception):
    pass


class ResearchDoesNotExist(Exception):
    pass


def make_suitable_name(name: str) -> str:
    return '-'.join(name.split())


def make_suitable_task_name(name: str) -> str:
    return '_'.join(name.split())


def make_suitable_research_dir(descr: str) -> str:
    return '-'.join([str(date.today()), make_suitable_name(descr)])


def get_task_full_name(task_number: int, task_name: str) -> str:
    return str(task_number) + '-' + make_suitable_task_name(task_name)


def split_task_dir(task_dir: str) -> (int, str):
    parsing_params = parse_by_named_regexp(r'^(?P<task_number>\d+)-(?P<task_name>\S+)', task_dir)
    if parsing_params is None:
        raise Exception("No task directory '{}' is found".format(task_dir))
    return int(parsing_params['task_number']), parsing_params['task_name']


def retrieve_trailing_float_from_task_dir(task_dir: str) -> float:
    matching = re.search(r'^(?P<task_number>\d+)-(?P<task_name>\S+)_(?P<float_left>\d+)\.(?P<float_right>\d+)',
                         task_dir)
    if matching is None:
        raise Exception('Incorrect task directory is given')
    return float('{}.{}'.format(matching.group('float_left'), matching.group('float_right')))


class CreateTaskEdge(Edge):
    def __init__(self, res, task_name_maker, predicate=dummy_predicate, remote=False):
        self._res = res
        self._task_name_maker = task_name_maker
        self._remote = remote
        super().__init__(predicate, Func(func=self.execute))

    def execute(self, data):
        task_name = self._task_name_maker(data)
        task_number = self._res.create_task(task_name)
        data['__WORKING_DIR__'] = self._res.get_task_path(task_number)
        if self._remote:
            data['__REMOTE_WORKING_DIR__'] = self._res.get_task_path(task_number, at_remote_host=True)


class CreateTaskGraph(Graph):
    def __init__(self, res, task_name_maker, array_keys_mapping=None, remote=False):
        s_init, s_term = self.create_branch(res, task_name_maker, array_keys_mapping=array_keys_mapping, remote=remote)
        super().__init__(s_init, s_term)

    @staticmethod
    def create_branch(res, task_name_maker, array_keys_mapping=None, remote=False):
        s_init = State('READY_FOR_TASK_CREATION', array_keys_mapping=array_keys_mapping)
        s_term = State('TASK_CREATED')
        s_init.connect_to(s_term, edge=CreateTaskEdge(res, task_name_maker=task_name_maker, remote=remote))
        return s_init, s_term