# Copyright (C) 2003  CAMP
# Please see the accompanying LICENSE file for further information.
"""K-point/spin combination-descriptors
This module contains classes for defining combinations of two indices:
* Index k for irreducible kpoints in the 1st Brillouin zone.
* Index s for spin up/down if spin-polarized (otherwise ignored).
"""
import numpy as np
from ase.units import Bohr
from ase.dft.kpoints import monkhorst_pack, get_monkhorst_pack_size_and_offset
from gpaw.symmetry import Symmetry
from gpaw.kpoint import KPoint
import gpaw.mpi as mpi
import _gpaw
class KPointDescriptor:
    """Descriptor-class for k-points."""
    def __init__(self, kpts, nspins=1, collinear=True):
        """Construct descriptor object for kpoint/spin combinations (ks-pair).
        Parameters
        ----------
        kpts: None, sequence of 3 ints, or (n,3) shaped ndarray
            Specification of the k-point grid. None=Gamma, list of
            ints=Monkhorst-Pack, ndarray=user specified.
        nspins: int
            Number of spins.
        Attributes
        ============  ======================================================
        ``N_c``       Number of k-points in the different directions.
        ``nspins``    Number of spins in total.
        ``mynspins``  Number of spins on this CPU.
        ``nibzkpts``  Number of irreducible kpoints in 1st Brillouin zone.
        ``nks``       Number of k-point/spin combinations in total.
        ``mynks``     Number of k-point/spin combinations on this CPU.
        ``gamma``     Boolean indicator for gamma point calculation.
        ``comm``      MPI-communicator for kpoint distribution.
        ============  ======================================================
        
        """
        if kpts is None:
            self.bzk_kc = np.zeros((1, 3))
            self.N_c = np.array((1, 1, 1), dtype=int)
            self.offset_c = np.zeros(3)
        elif isinstance(kpts[0], int):
            self.bzk_kc = monkhorst_pack(kpts)
            self.N_c = np.array(kpts, dtype=int)
            self.offset_c = np.zeros(3)
        else:
            self.bzk_kc = np.array(kpts, float)
            try:
                self.N_c, self.offset_c = \
                          get_monkhorst_pack_size_and_offset(self.bzk_kc)
            except ValueError:
                self.N_c = None
                self.offset_c = None
        self.collinear = collinear
        self.nspins = nspins
        self.nbzkpts = len(self.bzk_kc)
        
        # Gamma-point calculation?
        self.gamma = self.nbzkpts == 1 and not self.bzk_kc[0].any()
            
        self.set_symmetry(None, None, usesymm=None)
        self.set_communicator(mpi.serial_comm)
        if self.gamma:
            self.description = '1 k-point (Gamma)'
        else:
            self.description = '%d k-points' % self.nbzkpts
            if self.N_c is not None:
                self.description += (' (%d x %d x %d Monkhorst-Pack grid' %
                                     tuple(self.N_c))
                if self.offset_c.any():
                    self.description += ' + ['
                    for x in self.offset_c:
                        if x != 0 and abs(round(1 / x) - 1 / x) < 1e-12:
                            self.description += '1/%d,' % (1 / x)
                        else:
                            self.description += '%f,' % x
                    self.description = self.description[:-1] + ']'
                self.description += ')'
    def __len__(self):
        """Return number of k-point/spin combinations of local CPU."""
        
        return self.mynks
    def set_symmetry(self, atoms, setups, magmom_av=None,
                     usesymm=False, N_c=None, comm=None):
        """Create symmetry object and construct irreducible Brillouin zone.
        atoms: Atoms object
            Defines atom positions and types and also unit cell and
            boundary conditions.
        setups: instance of class Setups
            PAW setups for the atoms.
        magmom_av: ndarray
            Initial magnetic moments.
        usesymm: bool
            Symmetry flag.
        N_c: three int's or None
            If not None:  Check also symmetry of grid.
        """
        if atoms is not None:
            if (~atoms.pbc & self.bzk_kc.any(0)).any():
                raise ValueError('K-points can only be used with PBCs!')
            if magmom_av is None:
                magmom_av = np.zeros((len(atoms), 3))
                magmom_av[:, 2] = atoms.get_initial_magnetic_moments()
            magmom_av = magmom_av.round(decimals=3)  # round off
            id_a = zip(setups.id_a, *magmom_av.T)
            # Construct a Symmetry instance containing the identity operation
            # only
            self.symmetry = Symmetry(id_a, atoms.cell / Bohr, atoms.pbc)
        else:
            self.symmetry = None
        
        if self.gamma or usesymm is None:
            # Point group and time-reversal symmetry neglected
            self.weight_k = np.ones(self.nbzkpts) / self.nbzkpts
            self.ibzk_kc = self.bzk_kc.copy()
            self.sym_k = np.zeros(self.nbzkpts, int)
            self.time_reversal_k = np.zeros(self.nbzkpts, bool)
            self.bz2ibz_k = np.arange(self.nbzkpts)
            self.ibz2bz_k = np.arange(self.nbzkpts)
            self.bz2bz_ks = np.arange(self.nbzkpts)[:, np.newaxis]
        else:
            if usesymm:
                # Find symmetry operations of atoms
                self.symmetry.analyze(atoms.get_scaled_positions())
                
                if N_c is not None:
                    self.symmetry.prune_symmetries_grid(N_c)
            (self.ibzk_kc, self.weight_k,
             self.sym_k,
             self.time_reversal_k,
             self.bz2ibz_k,
             self.ibz2bz_k,
             self.bz2bz_ks) = self.symmetry.reduce(self.bzk_kc, comm)
            
        if setups is not None:
            setups.set_symmetry(self.symmetry)
        # Number of irreducible k-points and k-point/spin combinations.
        self.nibzkpts = len(self.ibzk_kc)
        if self.collinear:
            self.nks = self.nibzkpts * self.nspins
        else:
            self.nks = self.nibzkpts
        # Wrap k-points to 1. BZ:
        self.i1bzk_kc = self.ibzk_kc.copy()
        if atoms is not None:
            B_cv = 2.0 * np.pi * np.linalg.inv(atoms.cell / Bohr).T
            K_kv = np.dot(self.ibzk_kc, B_cv)
            N_xc = np.indices((3, 3, 3)).reshape((3, 27)).T - 1
            G_xv = np.dot(N_xc, B_cv)
            for k, K_v in enumerate(K_kv):
                x = ((G_xv - K_v)**2).sum(1).argmin()
                self.i1bzk_kc[k] -= N_xc[x]
        
    def set_communicator(self, comm):
        """Set k-point communicator."""
        # Ranks < self.rank0 have mynks0 k-point/spin combinations and
        # ranks >= self.rank0 have mynks0+1 k-point/spin combinations.
        mynks0, x = divmod(self.nks, comm.size)
        self.rank0 = comm.size - x
        self.comm = comm
        # My number and offset of k-point/spin combinations
        self.mynks, self.ks0 = self.get_count(), self.get_offset()
        if self.nspins == 2 and comm.size == 1:  # NCXXXXXXXX
            # Avoid duplicating k-points in local list of k-points.
            self.ibzk_qc = self.ibzk_kc.copy()
            self.i1bzk_qc = self.i1bzk_kc.copy()
        else:
            self.ibzk_qc = np.vstack((self.ibzk_kc,
                                      self.ibzk_kc))[self.get_slice()]
            self.i1bzk_qc = np.vstack((self.i1bzk_kc,
                                       self.i1bzk_kc))[self.get_slice()]
    def create_k_points(self, gd):
        """Return a list of KPoints."""
        sdisp_cd = gd.sdisp_cd
        kpt_u = []
        for ks in range(self.ks0, self.ks0 + self.mynks):
            s, k = divmod(ks, self.nibzkpts)
            q = (ks - self.ks0) % self.nibzkpts
            if self.collinear:
                weight = self.weight_k[k] * 2 / self.nspins
            else:
                weight = self.weight_k[k]
            if self.gamma:
                phase_cd = np.ones((3, 2), complex)
            else:
                phase_cd = np.exp(2j * np.pi *
                                  sdisp_cd * self.ibzk_kc[k, :, np.newaxis])
            kpt_u.append(KPoint(weight, s, k, q, phase_cd))
        return kpt_u
    def collect(self, a_ux, broadcast=True):
        """Collect distributed data to all."""
        if self.comm.rank == 0 or broadcast:
            xshape = a_ux.shape[1:]
            a_skx = np.empty((self.nspins, self.nibzkpts) + xshape, a_ux.dtype)
            a_Ux = a_skx.reshape((-1,) + xshape)
        else:
            a_skx = None
        if self.comm.rank > 0:
            self.comm.send(a_ux, 0)
        else:
            u1 = self.get_count(0)
            a_Ux[0:u1] = a_ux
            requests = []
            for rank in range(1, self.comm.size):
                u2 = u1 + self.get_count(rank)
                requests.append(self.comm.receive(a_Ux[u1:u2], rank,
                                                  block=False))
                u1 = u2
            assert u1 == len(a_Ux)
            self.comm.waitall(requests)
        
        if broadcast:
            self.comm.broadcast(a_Ux, 0)
        return a_skx
    def transform_wave_function(self, psit_G, k):
        """Transform wave function from IBZ to BZ.
        k is the index of the desired k-point in the full BZ.
        """
        
        s = self.sym_k[k]
        time_reversal = self.time_reversal_k[k]
        op_cc = np.linalg.inv(self.symmetry.op_scc[s]).round().astype(int)
        # Identity
        if (np.abs(op_cc - np.eye(3, dtype=int)) < 1e-10).all():
            if time_reversal:
                return psit_G.conj()
            else:
                return psit_G
        # General point group symmetry
        else:


Answer:             ik = self.bz2ibz_k[k]