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

""" bg_orbitals.py: 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, collect_screen_mpi_time
from bg_mpi_orbitals import orb_generator_master, orb_entanglement_main_par
from bg_print import print_orb_info, print_update

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

def orb_generator(molecule,dom,tup,l_limit,u_limit,k,level):
   #
   if (molecule['mpi_parallel']):
      #
      orb_generator_master(molecule,dom,tup,l_limit,u_limit,k,level)
   #
   else:
      #
      timer_mpi(molecule,'mpi_time_work_screen',k)
      #
      tmp_2 = []
      #
      for i in range(0,len(tup[k-1])):
         #
         # generate subset of all pairs within the parent tuple
         #
         parent_tup = tup[k-1][i]
         #
         tmp = list(list(comb) for comb in combinations(parent_tup,2))
         #
         mask = True
         #
         for j in range(0,len(tmp)):
            #
            # is the parent tuple still allowed?
            #
            if (not (set([tmp[j][1]]) < set(dom[-1][(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[-1]+1,(l_limit+u_limit)+1):
               #
               mask_2 = True
               #
               for l in parent_tup:
                  #
                  # is the new child tuple allowed?
                  #
                  if (not (set([m]) < set(dom[-1][(l-l_limit)-1]))):
                     #
                     mask_2 = False
                     #
                     break
               #
               if (mask_2):
                  #
                  # append the child tuple to the tup list
                  #
                  tmp_2.append(list(deepcopy(parent_tup)))
                  #
                  tmp_2[-1].append(m)
      #
      if (len(tmp_2) > 1): tmp_2.sort()
      #
      tup.append(np.array(tmp_2,dtype=np.int32))
      #
      del tmp
      del tmp_2
      #
      timer_mpi(molecule,'mpi_time_work_screen',k,True)
   #
   return tup

def orb_screening(molecule,l_limit,u_limit,order,level,calc_end=False):
   #
   if (order == 1):
      #
      # add singles contributions to orb_con list
      #
      orb_contributions(molecule,order,level,True)
      #
      # print orb info
      #
      if (molecule['debug']): print_orb_info(molecule,l_limit,u_limit,level)
      #
      # update domains
      #
      update_domains(molecule,l_limit,level,False)
   #
   else:
      #
      # set up entanglement and exclusion lists
      #
      orb_entanglement_main(molecule,l_limit,u_limit,order,level,calc_end)
      #
      timer_mpi(molecule,'mpi_time_work_screen',order)
      #
      orb_entanglement_arr(molecule,l_limit,u_limit,level)
      #
      orb_contributions(molecule,order,level)
      #
      if (calc_end):
         #
         if (molecule['mpi_parallel']):
            #
            collect_screen_mpi_time(molecule,order,True)
         #
         else:
            #
            timer_mpi(molecule,'mpi_time_work_screen',order,True)
      #
      else:
         #
         # print orb info
         #
         if (molecule['debug']): print_orb_info(molecule,l_limit,u_limit,level)
         #
         # construct exclusion list
         #
         orb_exclusion(molecule,l_limit,order,level)
         #
         # update domains
         #
         if (order == 2):
            #
            update_domains(molecule,l_limit,level,False)
         #
         else:
            #
            update_domains(molecule,l_limit,level)
         #
         # print domain updates
         #
         print_update(molecule,l_limit,u_limit,level)
         #
         # update threshold and restart frequency
         #
         update_thres_and_rst_freq(molecule,level)
         #
         if (molecule['mpi_parallel']):
            #
            collect_screen_mpi_time(molecule,order)
         #
         else:
            #
            timer_mpi(molecule,'mpi_time_work_screen',order,True)
   #
   return molecule

def orb_entanglement_main(molecule,l_limit,u_limit,order,level,calc_end):
   #
   if (molecule['mpi_parallel']):
      #
      orb_entanglement_main_par(molecule,l_limit,u_limit,order,level,calc_end)
   #
   else:
      #
      timer_mpi(molecule,'mpi_time_work_screen',order)
      #
      orb = molecule['prim_orb_ent']
      #
      orb.append(np.zeros([u_limit,u_limit],dtype=np.float64))
      #
      for l in range(0,len(molecule['prim_tuple'][order-1])):
         #
         tup = molecule['prim_tuple'][order-1]           
         e_inc = molecule['prim_energy_inc'][order-1]
         ldx = l
         # 
         for i in range(l_limit,l_limit+u_limit):
            #
            for j in range(i+1,l_limit+u_limit):
               #
               # add up contributions from the correlation between orbs i and j at current order
               #
               if (set([i+1,j+1]) <= set(tup[ldx])):
                  #
                  orb[-1][i-l_limit,j-l_limit] += e_inc[ldx]
                  orb[-1][j-l_limit,i-l_limit] = orb[-1][i-l_limit,j-l_limit]
      #
      timer_mpi(molecule,'mpi_time_work_screen',order,True)
   #
   return molecule
      
def orb_entanglement_arr(molecule,l_limit,u_limit,level):
   #
   if (level == 'MACRO'):
      #
      orb = molecule['prim_orb_ent']
      orb_arr = molecule['prim_orb_arr']
   #
   # write orbital entanglement matrix
   #
   orb_arr.append(np.empty([molecule['u_limit'],molecule['u_limit']],dtype=np.float64))
   #
   tmp = np.zeros(molecule['u_limit'],dtype=np.float64)
   #
   for i in range(0,len(orb[-1])):
      #
      tmp.fill(0.0)
      #
      # sum up contributions between orbital 'i' and all other orbitals
      #
      for k in range(0,len(orb)):
         #
         for j in range(0,len(orb[-1][i])):
            #
            tmp[j] += orb[k][i,j]
      #
      # account for what has already been screened away
      #
      for j in range(0,len(orb[-1][i])):
         #
         if (orb[-1][i,j] == 0.0): tmp[j] = 0.0 
      #
      # make result relative to the maximum contribution
      #
      for j in range(0,len(orb[-1][i])):
         #
         if (tmp[j] == 0.0):
            #
            orb_arr[-1][i,j] = 0.0
         #
         else:
            #
            if (np.abs(np.max(tmp)) > np.abs(np.min(tmp))):
               #
               orb_arr[-1][i,j] = tmp[j]/np.max(tmp)
            #
            else:
               #
               orb_arr[-1][i,j] = tmp[j]/np.min(tmp)
   #
   del tmp
   #
   return molecule

def orb_contributions(molecule,order,level,singles=False):
   #
   if (level == 'MACRO'):
      #
      orb = molecule['prim_orb_ent']
      orb_con_abs = molecule['prim_orb_con_abs']
      orb_con_rel = molecule['prim_orb_con_rel']
   #
   if (singles):
      #
      e_inc = molecule['prim_energy_inc'][order-1]
      #
      # total orbital contribution
      #
      orb_con_abs.append([])
      orb_con_rel.append([])
      #
      if (((molecule['exp'] == 'occ') or (molecule['exp'] == 'comb-ov')) and (molecule['frozen'])):
         #
         for _ in range(0,molecule['ncore']):
            #
            orb_con_abs[-1].append(0.0)
      #
      for i in range(0,len(e_inc)):
         #
         orb_con_abs[-1].append(e_inc[i])
      #
      for i in range(0,len(orb_con_abs[-1])):
         #
         orb_con_rel[-1].append(abs(orb_con_abs[-1][i])/np.abs(np.sum(e_inc)))
   #
   else:
      #
      orb_con_abs.append([])
      orb_con_rel.append([])
      #
      # total orbital contribution
      #
      for i in range(0,len(orb[-1])):
         #
         orb_con_abs[-1].append(orb_con_abs[-2][i]+np.sum(orb[-1][i]))
      #
      for i in range(0,len(orb_con_abs[-1])):
         #
         if (orb_con_abs[-1][i] == 0.0):
            #
            orb_con_rel[-1].append(0.0)
         #
         else:
            #
            orb_con_rel[-1].append(orb_con_abs[-1][i]/sum(orb_con_abs[-1]))
   #
   return molecule

def orb_exclusion(molecule,l_limit,order,level):
   #
   if (level == 'MACRO'):
      #
      orb = molecule['prim_orb_ent']
      orb_arr = molecule['prim_orb_arr']
      thres = (molecule['prim_thres']/100.0)
   #
   molecule['excl_list'].append([])
   #
   # screening in individual domains based on orbital entanglement 
   #
   for i in range(0,len(orb_arr[-1])):
      #
      molecule['excl_list'][-1].append([])
      #
      if (np.sum(orb_arr[-1][i]) > 0.0):
         #
         for j in range(0,len(orb_arr[-1][i])):
            #
            if ((np.abs(orb_arr[-1][i,j]) < thres) and (orb[-1][i,j] != 0.0)): molecule['excl_list'][-1][i].append((j+l_limit)+1)
   #
   for i in range(0,len(molecule['excl_list'][-1])):
      #
      molecule['excl_list'][-1][i].sort()
   #
   return molecule

def update_domains(molecule,l_limit,level,screen=True):
   #
   if (level == 'MACRO'):
      #
      dom = molecule['prim_domain']
   #
   dom.append([])
   #
   for l in range(0,len(dom[0])):
      #
      dom[-1].append(list(dom[-2][l]))
   #
   if (screen):
      #
      for i in range(0,len(molecule['excl_list'][-1])):
         #
         for j in range(0,len(molecule['excl_list'][-1][i])):
            #
            if (molecule['excl_list'][-1][i][j] in molecule['excl_list'][-2][i]):
               #
               dom[-1][i].remove(molecule['excl_list'][-1][i][j])
               dom[-1][(molecule['excl_list'][-1][i][j]-l_limit)-1].remove((i+l_limit)+1)
               #
               if ((i+l_limit)+1 in molecule['excl_list'][-1][(molecule['excl_list'][-1][i][j]-l_limit)-1]):
                  #
                  molecule['excl_list'][-1][(molecule['excl_list'][-1][i][j]-l_limit)-1].remove((i+l_limit)+1)
   #
   return molecule

def init_domains(molecule):
   #
   molecule['occ_domain'] = []
   molecule['virt_domain'] = []
   #
   for i in range(0,molecule['nocc']):
      #
      molecule['occ_domain'].append(list(range(1,molecule['nocc']+1)))
      #
      molecule['occ_domain'][i].pop(i)
   #
   if (molecule['frozen']):
      #
      for i in range(0,molecule['ncore']):
         #
         molecule['occ_domain'][i][:] = []
      #
      for j in range(molecule['ncore'],molecule['nocc']):
         #
         for _ in range(0,molecule['ncore']):
            #
            molecule['occ_domain'][j].pop(0)
   #
   for i in range(0,molecule['nvirt']):
      #
      molecule['virt_domain'].append(list(range(molecule['nocc']+1,(molecule['nocc']+molecule['nvirt'])+1)))
      #
      molecule['virt_domain'][i].pop(i)
   #
   return molecule

def update_thres_and_rst_freq(molecule,level):
   #
   # update threshold by doubling it
   #
   if (level == 'MACRO'):
      #
      molecule['prim_thres'] += molecule['prim_thres']
   #
   # update restart frequency by halving it
   #
   molecule['rst_freq'] /= 2.
   #
   return molecule