#!/usr/bin/env python
# -*- coding: utf-8 -*

""" bg_mpi_orbitals.py: MPI orbital-related routines for Bethe-Goldstone correlation calculations."""

import numpy as np
from mpi4py import MPI
from itertools import combinations
from copy import deepcopy

from bg_mpi_time import timer_mpi

__author__ = 'Dr. Janus Juul Eriksen, JGU Mainz'
__copyright__ = 'Copyright 2017'
__credits__ = ['Prof. Juergen Gauss', 'Dr. Filippo Lipparini']
__license__ = '???'
__version__ = '0.4'
__maintainer__ = 'Dr. Janus Juul Eriksen'
__email__ = 'jeriksen@uni-mainz.de'
__status__ = 'Development'

def orb_generator_master(molecule,dom,tup,l_limit,u_limit,k,level):
   #
   #  ---  master routine
   #
   if ((k <= 2)):
      #
      # wake up slaves
      #
      msg = {'task': 'bcast_tuples', 'order': k}
      #
      molecule['mpi_comm'].bcast(msg,root=0)
      #
      timer_mpi(molecule,'mpi_time_work_init',k)
      #
      if (((molecule['exp'] == 'occ') or (molecule['exp'] == 'comb-ov')) and molecule['frozen']):
         #
         start = molecule['ncore']
      #
      else:
         #
         start = 0
      #
      if (k == 1):
         #
         # all singles contributions
         #
         tmp = []
         #
         for i in range(start,len(dom)):
            #
            tmp.append([(i+l_limit)+1])
         #
         tup.append(np.array(tmp,dtype=np.int))
      #
      elif (k == 2):
         #
         # generate all possible (unique) pairs
         #
         tmp = list(list(comb) for comb in combinations(range(start+(1+l_limit),(l_limit+u_limit)+1),2))
         #
         tup.append(np.array(tmp,dtype=np.int))
      #
      # bcast final dummy message to collect idle slave time
      #
      timer_mpi(molecule,'mpi_time_comm_init',k)
      #
      final_msg = {'done': None}
      #
      molecule['mpi_comm'].bcast(final_msg,root=0)
      #
      # bcast total number of tuples
      #
      tup_info = {'tot_tup': len(tup[k-1])}
      #
      molecule['mpi_comm'].bcast(tup_info,root=0)
      #
      # bcast the tuples
      #
      molecule['mpi_comm'].Bcast([tup[k-1],MPI.INT],root=0)
      #
      timer_mpi(molecule,'mpi_time_comm_init',k,True)
   #
   else:
      #
      # init job_info dictionary
      #
      job_info = {}
      #
      # number of slaves
      #
      num_slaves = molecule['mpi_size'] - 1
      #
      # number of available slaves
      #
      slaves_avail = num_slaves
      #
      # define mpi message tags
      #
      tags = enum('ready','comm','done','exit','start')
      #
      # init job index
      #
      i = 0
      #
      # wake up slaves
      #
      timer_mpi(molecule,'mpi_time_comm_init',k)
      #
      msg = {'task': 'orb_generator_par', 'l_limit': l_limit, 'u_limit': u_limit, 'order': k, 'level': level}
      #
      molecule['mpi_comm'].bcast(msg,root=0)
      #
      # bcast orbital domains and lower/upper limits
      #
      dom_info = {'dom': dom}
      #
      molecule['mpi_comm'].bcast(dom_info,root=0)
      #
      timer_mpi(molecule,'mpi_time_work_init',k)
      #
      # init tmp list
      #
      tmp = []
      #
      # init parent_tup
      #
      if (level == 'MACRO'):
         #
         parent_tup = tup[k-2]
      #
      elif (level == 'CORRE'):
         #
         if (k == molecule['min_corr_order']):
            #
            parent_tup = molecule['prim_tuple'][k-2]
         #
         else:
            #
            parent_tup = np.vstack((tup[k-2],molecule['prim_tuple'][k-2]))
      #
      while (slaves_avail >= 1):
         #
         # receive data dict
         #
         timer_mpi(molecule,'mpi_time_idle_init',k)
         #
         data = molecule['mpi_comm'].recv(source=MPI.ANY_SOURCE,tag=MPI.ANY_TAG,status=molecule['mpi_stat'])
         #
         timer_mpi(molecule,'mpi_time_work_init',k)
         #
         # probe for source
         #
         source = molecule['mpi_stat'].Get_source()
         #
         # probe for tag
         #
         tag = molecule['mpi_stat'].Get_tag()
         #
         if (tag == tags.ready):
            #
            if (i <= (len(parent_tup)-1)):
               #
               # send parent tuple
               #
               timer_mpi(molecule,'mpi_time_comm_init',k)
               #
               # store index
               #
               job_info['index'] = i
               #
               molecule['mpi_comm'].send(job_info,dest=source,tag=tags.start)
               #
               timer_mpi(molecule,'mpi_time_work_init',k)
               #
               # increment job index
               #
               i += 1
            #
            else:
               #
               timer_mpi(molecule,'mpi_time_comm_init',k)
               #
               molecule['mpi_comm'].send(None,dest=source,tag=tags.exit)
               #
               timer_mpi(molecule,'mpi_time_work_init',k)
         #
         elif (tag == tags.comm):
            #
            # receive child_tup for source
            #
            timer_mpi(molecule,'mpi_time_comm_init',k)
            #
            data = molecule['mpi_comm'].recv(source=source,tag=tags.done,status=molecule['mpi_stat'])
            #
            # write tmp child tuple list
            #
            timer_mpi(molecule,'mpi_time_work_init',k)
            #
            tmp += data['child_tup'] 
         #
         elif (tag == tags.exit):
            #
            slaves_avail -= 1
      #
      # bcast final dummy message to collect idle slave time
      #
      timer_mpi(molecule,'mpi_time_idle_init',k)
      #
      final_msg = {'done': None}
      #
      molecule['mpi_comm'].bcast(final_msg,root=0)
      #
      # finally we sort the tuples
      #
      timer_mpi(molecule,'mpi_time_work_init',k)
      #
      if (len(tmp) >= 1): tmp.sort()
      #
      # append tup[k-1] with numpy array of tmp list
      #
      tup.append(np.array(tmp,dtype=np.int))
      #
      # bcast total number of tuples
      #
      timer_mpi(molecule,'mpi_time_comm_init',k)
      #
      tup_info = {'tot_tup': len(tup[k-1])}
      #
      molecule['mpi_comm'].bcast(tup_info,root=0)
      #
      # bcast the tuples
      #
      molecule['mpi_comm'].Bcast([tup[k-1],MPI.INT],root=0)
      #
      timer_mpi(molecule,'mpi_time_comm_init',k,True)
      #
      dom_info.clear()
      #
      del parent_tup
   #
   tup_info.clear()
   final_msg.clear()
   #
   del tmp
   #
   return tup

def orb_generator_slave(molecule,dom,tup,l_limit,u_limit,k,level):
   #
   #  ---  slave routine
   #
   # define mpi message tags
   #
   tags = enum('ready','comm','done','exit','start')
   #
   # prepare msg dict
   #
   msg = {'done': None}
   #
   # init data dict
   #
   data = {'child_tup': []}
   #
   # init tmp lists
   #
   tmp = []
   #
   timer_mpi(molecule,'mpi_time_work_init',k)
   #
   # init parent_tup
   #
   if (level == 'MACRO'):
      #
      parent_tup = tup[k-2]
   #
   elif (level == 'CORRE'):
      #
      if (k == molecule['min_corr_order']):
         #
         parent_tup = molecule['prim_tuple'][k-2]
      #
      else:
         #
         parent_tup = np.vstack((tup[k-2],molecule['prim_tuple'][k-2]))
   #
   while True:
      #
      # ready for task
      #
      timer_mpi(molecule,'mpi_time_comm_init',k)
      #
      molecule['mpi_comm'].send(None,dest=0,tag=tags.ready)
      #
      # receive parent tuple index
      #
      job_info = molecule['mpi_comm'].recv(source=0,tag=MPI.ANY_SOURCE,status=molecule['mpi_stat'])
      #
      timer_mpi(molecule,'mpi_time_work_init',k)
      #
      # recover tag
      #
      tag = molecule['mpi_stat'].Get_tag()
      #
      # do job or break out (exit)
      #
      if (tag == tags.start):
         #
         data['child_tup'][:] = []
         #
         # generate subset of all pairs within the parent tuple
         #
         tmp = list(list(comb) for comb in combinations(parent_tup[job_info['index']],2))
         #
         mask = True
         #
         for j in range(0,len(tmp)):
            #
            # is the parent tuple still allowed?
            #
            if (not (set([tmp[j][1]]) < set(dom[(tmp[j][0]-l_limit)-1]))):
               #
               mask = False
               #
               break
         #
         if (mask):
            #
            # loop through possible orbitals to augment the parent tuple with
            #
            for m in range(parent_tup[job_info['index']][-1]+1,(l_limit+u_limit)+1):
               #
               mask_2 = True
               #
               for l in parent_tup[job_info['index']]:
                  #
                  # is the new child tuple allowed?
                  #
                  if (not (set([m]) < set(dom[(l-l_limit)-1]))):
                     #
                     mask_2 = False
                     #
                     break
               #
               if (mask_2):
                  #
                  # append the child tuple to the tup list
                  #
                  data['child_tup'].append(list(deepcopy(parent_tup[job_info['index']])))
                  #
                  data['child_tup'][-1].append(m)
                  #
                  # check whether this tuple has already been accounted for in the primary expansion
                  #
                  if ((level == 'CORRE') and (np.equal(data['child_tup'][-1],molecule['prim_tuple'][k-1]).all(axis=1).any())):
                     #
                     data['child_tup'].pop(-1)
         #
         timer_mpi(molecule,'mpi_time_comm_init',k)
         #
         # prepare master for communication (to distinguish between idle and comm time on master)
         #
         molecule['mpi_comm'].send(msg,dest=0,tag=tags.comm)
         #
         # send child tuple back to master
         #
         molecule['mpi_comm'].send(data,dest=0,tag=tags.done)
         #
         timer_mpi(molecule,'mpi_time_work_init',k)
      #
      elif (tag == tags.exit):
         #
         break
   #
   # exit
   #
   timer_mpi(molecule,'mpi_time_comm_init',k)
   #
   molecule['mpi_comm'].send(None,dest=0,tag=tags.exit)
   #
   timer_mpi(molecule,'mpi_time_idle_init',k)
   #
   final_msg = MPI.COMM_WORLD.bcast(None,root=0)
   #
   # receive the total number of tuples
   #
   timer_mpi(molecule,'mpi_time_comm_init',k)
   #
   tup_info = MPI.COMM_WORLD.bcast(None,root=0)
   #
   # init tup[k-1]
   #
   timer_mpi(molecule,'mpi_time_work_init',k)
   #
   tup.append(np.empty([tup_info['tot_tup'],k],dtype=np.int))
   #
   # receive the tuples
   #
   timer_mpi(molecule,'mpi_time_comm_init',k)
   #
   MPI.COMM_WORLD.Bcast([tup[k-1],MPI.INT],root=0)
   #
   timer_mpi(molecule,'mpi_time_comm_init',k,True)
   #
   data.clear()
   msg.clear()
   final_msg.clear()
   tup_info.clear()
   #
   del tmp
   del parent_tup
   #
   return molecule

def bcast_dom(molecule,k):
   #
   #  ---  slave routine
   #
   timer_mpi(molecule,'mpi_time_comm_init',k)
   #
   molecule['dom_info'] = MPI.COMM_WORLD.bcast(None,root=0)
   #
   timer_mpi(molecule,'mpi_time_comm_init',k,True)
   #
   return molecule

def bcast_tuples(molecule,k):
   #
   #  ---  slave routine
   #
   timer_mpi(molecule,'mpi_time_idle_init',k)
   #
   final_msg = MPI.COMM_WORLD.bcast(None,root=0)
   #
   timer_mpi(molecule,'mpi_time_comm_init',k)
   #
   # receive the total number of tuples
   #
   tup_info = MPI.COMM_WORLD.bcast(None,root=0)
   #
   # init tup[k-1]
   #
   timer_mpi(molecule,'mpi_time_work_init',k)
   #
   molecule['prim_tuple'].append(np.empty([tup_info['tot_tup'],k],dtype=np.int))
   #
   # receive the tuples
   #
   timer_mpi(molecule,'mpi_time_comm_init',k)
   #
   MPI.COMM_WORLD.Bcast([molecule['prim_tuple'][k-1],MPI.INT],root=0)
   #
   timer_mpi(molecule,'mpi_time_comm_init',k,True)
   #
   tup_info.clear()
   final_msg.clear()
   #
   return molecule

def enum(*sequential,**named):
   #
   # hardcoded enums
   #
   enums = dict(zip(sequential,range(len(sequential))),**named)
   #
   return type('Enum',(), enums)