#
#    ICRAR - International Centre for Radio Astronomy Research
#    (c) UWA - The University of Western Australia, 2019
#    Copyright by UWA (in the framework of the ICRAR)
#    All rights reserved
#
#    This library is free software; you can redistribute it and/or
#    modify it under the terms of the GNU Lesser General Public
#    License as published by the Free Software Foundation; either
#    version 2.1 of the License, or (at your option) any later version.
#
#    This library is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#    Lesser General Public License for more details.
#
#    You should have received a copy of the GNU Lesser General Public
#    License along with this library; if not, write to the Free Software
#    Foundation, Inc., 59 Temple Place, Suite 330, Boston,
#    MA 02111-1307  USA
#
import collections
import logging
import os
import pickle
import time

from . import slurm_utils


logger = logging.getLogger(__name__)

class Remote(object):

    def __init__(self, options, my_ip):
        self.options = options
        self.my_ip = my_ip
        self.num_islands = options.num_islands
        self.run_proxy = options.monitor_host is not None
        self.rank = None
        self.size = None
        self.sorted_peers = None

    def _set_world(self, rank, size, sorted_peers):
        """Called by subclasses"""
        self.rank = rank
        self.size = size
        self.sorted_peers = sorted_peers
        if len(set(self.sorted_peers)) != self.size:
            raise RuntimeError('More than one task started per node, cannot continue')
        nm_range = self._nm_range()
        if nm_range[0] == nm_range[1]:
            raise RuntimeError('No nodes left for Node Managers with the requested setup')

    def _dim_range(self):
        first = 0
        if self.num_islands > 1:
            first = 1
        if self.run_proxy:
            first += 1
        return first, first + self.num_islands

    def _nm_range(self):
        first = 1
        if self.run_proxy:
            first = 2
        if self.num_islands > 1:
            first += self.num_islands
        return first, self.size

    def _dim_nms(self, pg):
        dim_nodes = collections.defaultdict(set)
        ranks = {ip: rank for rank, ip in enumerate(self.sorted_peers)}
        for drop in pg:
            nm, dim = drop['node'], drop['island']
            dim_nodes[dim].add(nm)
        for dim in self.dim_ips:
            r = ranks[dim]
            nms = list(dim_nodes[dim])
            logger.debug("Nodes for DIM %s (rank %d): %r", dim, r, nms)
            yield r, nms

    @property
    def is_highest_level_manager(self):
        return self.rank == 0

    @property
    def hl_mgr_ip(self):
        return self.sorted_peers[0]

    @property
    def is_dim(self):
        return self.rank in range(*self._dim_range())

    @property
    def dim_ips(self):
        start, end = self._dim_range()
        return self.sorted_peers[start:end]

    @property
    def is_nm(self):
        return self.rank in range(*self._nm_range())

    @property
    def nm_ips(self):
        start, end = self._nm_range()
        return self.sorted_peers[start:end]

    @property
    def is_proxy(self):
        return self.run_proxy and self.rank == 1

    @property
    def proxy_ip(self):
        if not self.run_proxy:
            raise RuntimeError('proxy_ip requested when not running with proxy enabled')
        return self.sorted_peers[1]


class MPIRemote(Remote):

    def __init__(self, options, my_ip):
        super(MPIRemote, self).__init__(options, my_ip)
        from mpi4py import MPI  # @UnresolvedImport
        self.comm = MPI.COMM_WORLD  # @UndefinedVariable
        self._set_world(self.comm.Get_rank(), self.comm.Get_size(),
                        self.comm.allgather(self.my_ip))

    def send_dim_nodes(self, pg):
        for dim_rank, nms in self._dim_nms(pg):
            self.comm.send(nms, dest=dim_rank)

    def recv_dim_nodes(self):
        return self.comm.recv(source=0)


class SlurmRemote(Remote):

    def __init__(self, options, my_ip):
        super(SlurmRemote, self).__init__(options, my_ip)
        self._set_world(
            int(os.environ['SLURM_PROCID']), int(os.environ['SLURM_NTASKS']),
            slurm_utils.list_as_string(os.environ['SLURM_NODELIST']))

    def send_dim_nodes(self, pg):
        basedir = self.options.log_dir
        for dim_rank, nms in self._dim_nms(pg):
            fname = 'dim_%d_nodes.pickle' % dim_rank
            stage1_fname = os.path.join(basedir, '.%s' % fname)
            stage2_fname = os.path.join(basedir, '%s' % fname)
            with open(stage1_fname, 'wb') as f:
                pickle.dump(nms, f)
            os.rename(stage1_fname, stage2_fname)

    def recv_dim_nodes(self):
        fname = os.path.join(self.options.log_dir, 'dim_%d_nodes.pickle' % self.rank)
        while not os.path.exists(fname):
            time.sleep(1)
        with open(fname, 'rb') as f:
            return pickle.load(f)