"""Define the PETSc Transfer class."""
from __future__ import division

import numpy as np
from petsc4py import PETSc
from six import iteritems, itervalues
from itertools import product, chain
from collections import defaultdict

from openmdao.vectors.transfer import Transfer
from openmdao.vectors.default_transfer import DefaultTransfer, _merge
from openmdao.vectors.vector import INT_DTYPE
from openmdao.utils.mpi import MPI
from openmdao.utils.array_utils import convert_neg, _flatten_src_indices

_empty_idx_array = np.array([], dtype=INT_DTYPE)


class PETScTransfer(DefaultTransfer):
    """
    PETSc Transfer implementation for running in parallel.

    Attributes
    ----------
    _scatter : method
        Method that performs a PETSc scatter.
    _transfer : method
        Method that performs either a normal transfer or a multi-transfer.
    """

    def __init__(self, in_vec, out_vec, in_inds, out_inds, comm):
        """
        Initialize all attributes.

        Parameters
        ----------
        in_vec : <Vector>
            pointer to the input vector.
        out_vec : <Vector>
            pointer to the output vector.
        in_inds : int ndarray
            input indices for the transfer.
        out_inds : int ndarray
            output indices for the transfer.
        comm : MPI.Comm or <FakeComm>
            communicator of the system that owns this transfer.
        """
        super(PETScTransfer, self).__init__(in_vec, out_vec, in_inds, out_inds, comm)
        in_indexset = PETSc.IS().createGeneral(self._in_inds, comm=self._comm)
        out_indexset = PETSc.IS().createGeneral(self._out_inds, comm=self._comm)

        self._scatter = PETSc.Scatter().create(out_vec._petsc, out_indexset, in_vec._petsc,
                                               in_indexset).scatter

        if in_vec._ncol > 1:
            self._transfer = self._multi_transfer

    @staticmethod
    def _setup_transfers(group, recurse=True):
        """
        Compute all transfers that are owned by our parent group.

        Parameters
        ----------
        group : <Group>
            Parent group.
        recurse : bool
            Whether to call this method in subsystems.
        """
        rev = group._mode != 'fwd'

        if recurse:
            for subsys in group._subgroups_myproc:
                subsys._setup_transfers(recurse)

        abs2meta = group._var_abs2meta
        allprocs_abs2meta = group._var_allprocs_abs2meta
        myproc = group.comm.rank

        transfers = group._transfers = {}
        vectors = group._vectors
        offsets = group._get_var_offsets()

        vec_names = group._lin_rel_vec_name_list if group._use_derivatives else group._vec_names

        mypathlen = len(group.pathname + '.' if group.pathname else '')
        sub_inds = group._subsystems_inds

        for vec_name in vec_names:
            relvars, _ = group._relevant[vec_name]['@all']

            # Initialize empty lists for the transfer indices
            nsub_allprocs = len(group._subsystems_allprocs)
            xfer_in = []
            xfer_out = []
            fwd_xfer_in = [[] for s in group._subsystems_allprocs]
            fwd_xfer_out = [[] for s in group._subsystems_allprocs]
            if rev:
                rev_xfer_in = [[] for s in group._subsystems_allprocs]
                rev_xfer_out = [[] for s in group._subsystems_allprocs]

            allprocs_abs2idx = group._var_allprocs_abs2idx[vec_name]
            sizes_in = group._var_sizes[vec_name]['input']
            sizes_out = group._var_sizes[vec_name]['output']
            offsets_in = offsets[vec_name]['input']
            offsets_out = offsets[vec_name]['output']

            # Loop through all connections owned by this system
            for abs_in, abs_out in iteritems(group._conn_abs_in2out):
                if abs_out not in relvars['output'] or abs_in not in relvars['input']:
                    continue

                # Only continue if the input exists on this processor
                if abs_in in abs2meta:

                    # Get meta
                    meta_in = abs2meta[abs_in]
                    meta_out = allprocs_abs2meta[abs_out]

                    idx_in = allprocs_abs2idx[abs_in]
                    idx_out = allprocs_abs2idx[abs_out]

                    # Read in and process src_indices
                    src_indices = meta_in['src_indices']
                    if src_indices is None:
                        owner = group._owning_rank[abs_out]
                        # if the input is larger than the output on a single proc, we have
                        # to just loop over the procs in the same way we do when src_indices
                        # is defined.
                        if meta_in['size'] > sizes_out[owner, idx_out]:
                            src_indices = np.arange(meta_in['size'], dtype=INT_DTYPE)
                    elif src_indices.ndim == 1:
                        src_indices = convert_neg(src_indices, meta_out['global_size'])
                    else:
                        src_indices = _flatten_src_indices(src_indices, meta_in['shape'],
                                                           meta_out['global_shape'],
                                                           meta_out['global_size'])

                    # 1. Compute the output indices
                    if src_indices is None:
                        rank = myproc if abs_out in abs2meta else owner
                        offset = offsets_out[rank, idx_out]
                        output_inds = np.arange(offset, offset + meta_in['size'], dtype=INT_DTYPE)
                    else:
                        output_inds = np.zeros(src_indices.size, INT_DTYPE)
                        start = end = 0
                        for iproc in range(group.comm.size):
                            end += sizes_out[iproc, idx_out]
                            if start == end:
                                continue

                            # The part of src on iproc
                            on_iproc = np.logical_and(start <= src_indices, src_indices < end)

                            if np.any(on_iproc):
                                # This converts from iproc-then-ivar to ivar-then-iproc ordering
                                # Subtract off part of previous procs
                                # Then add all variables on previous procs
                                # Then all previous variables on this proc
                                # - np.sum(out_sizes[:iproc, idx_out])
                                # + np.sum(out_sizes[:iproc, :])
                                # + np.sum(out_sizes[iproc, :idx_out])
                                # + inds
                                offset = offsets_out[iproc, idx_out] - start
                                output_inds[on_iproc] = src_indices[on_iproc] + offset

                            start = end

                    # 2. Compute the input indices
                    input_inds = np.arange(offsets_in[myproc, idx_in],
                                           offsets_in[myproc, idx_in] +
                                           sizes_in[myproc, idx_in], dtype=INT_DTYPE)

                    # Now the indices are ready - input_inds, output_inds
                    sub_in = abs_in[mypathlen:].split('.', 1)[0]
                    isub = sub_inds[sub_in]
                    fwd_xfer_in[isub].append(input_inds)
                    fwd_xfer_out[isub].append(output_inds)
                    if rev:
                        sub_out = abs_out[mypathlen:].split('.', 1)[0]
                        isub = sub_inds[sub_out]
                        rev_xfer_in[isub].append(input_inds)
                        rev_xfer_out[isub].append(output_inds)

            transfers[vec_name] = {}

            for isub in range(nsub_allprocs):
                fwd_xfer_in[isub] = _merge(fwd_xfer_in[isub])
                fwd_xfer_out[isub] = _merge(fwd_xfer_out[isub])
                if rev:
                    rev_xfer_in[isub] = _merge(rev_xfer_in[isub])
                    rev_xfer_out[isub] = _merge(rev_xfer_out[isub])

            xfer_in = np.concatenate(fwd_xfer_in)
            xfer_out = np.concatenate(fwd_xfer_out)

            out_vec = vectors['output'][vec_name]

            xfer_all = PETScTransfer(vectors['input'][vec_name], out_vec,
                                     xfer_in, xfer_out, group.comm)

            transfers[vec_name]['fwd', None] = xfer_all
            if rev:
                transfers[vec_name]['rev', None] = xfer_all

            for isub in range(nsub_allprocs):
                transfers[vec_name]['fwd', isub] = PETScTransfer(
                    vectors['input'][vec_name], vectors['output'][vec_name],
                    fwd_xfer_in[isub], fwd_xfer_out[isub], group.comm)
                if rev:
                    transfers[vec_name]['rev', isub] = PETScTransfer(
                        vectors['input'][vec_name], vectors['output'][vec_name],
                        rev_xfer_in[isub], rev_xfer_out[isub], group.comm)

        if group._use_derivatives:
            transfers['nonlinear'] = transfers['linear']

    def _transfer(self, in_vec, out_vec, mode='fwd'):
        """
        Perform transfer.

        Parameters
        ----------
        in_vec : <Vector>
            pointer to the input vector.
        out_vec : <Vector>
            pointer to the output vector.
        mode : str
            'fwd' or 'rev'.
        """
        flag = False
        if mode == 'rev':
            flag = True
            in_vec, out_vec = out_vec, in_vec

        in_petsc = in_vec._petsc
        out_petsc = out_vec._petsc

        # For Complex Step, need to disassemble real and imag parts, transfer them separately,
        # then reassemble them.
        if in_vec._under_complex_step and out_vec._alloc_complex:

            # Real
            in_petsc.array = in_vec._data.real
            out_petsc.array = out_vec._data.real
            self._scatter(out_petsc, in_petsc, addv=flag, mode=flag)

            # Imaginary
            in_petsc_imag = in_vec._imag_petsc
            out_petsc_imag = out_vec._imag_petsc
            in_petsc_imag.array = in_vec._data.imag
            out_petsc_imag.array = out_vec._data.imag
            self._scatter(out_petsc_imag, in_petsc_imag, addv=flag, mode=flag)

            in_vec._data[:] = in_petsc.array + in_petsc_imag.array * 1j

        else:

            # Anything that has been allocated complex requires an additional step because
            # the petsc vector does not directly reference the _data.

            if in_vec._alloc_complex:
                in_petsc.array = in_vec._data

            if out_vec._alloc_complex:
                out_petsc.array = out_vec._data

            self._scatter(out_petsc, in_petsc, addv=flag, mode=flag)

            if in_vec._alloc_complex:
                in_vec._data[:] = in_petsc.array

    def _multi_transfer(self, in_vec, out_vec, mode='fwd'):
        """
        Perform transfer.

        Parameters
        ----------
        in_vec : <Vector>
            pointer to the input vector.
        out_vec : <Vector>
            pointer to the output vector.
        mode : str
            'fwd' or 'rev'.
        """
        if mode == 'fwd':
            in_petsc = in_vec._petsc
            out_petsc = out_vec._petsc

            # For Complex Step, need to disassemble real and imag parts, transfer them separately,
            # then reassemble them.
            if in_vec._under_complex_step and out_vec._alloc_complex:
                in_petsc_imag = in_vec._imag_petsc
                out_petsc_imag = out_vec._imag_petsc
                for i in range(in_vec._ncol):

                    # Real
                    in_petsc.array = in_vec._data[:, i].real
                    out_petsc.array = out_vec._data[:, i].real
                    self._scatter(out_petsc, in_petsc, addv=False, mode=False)

                    # Imaginary
                    in_petsc_imag.array = in_vec._data[:, i].imag
                    out_petsc_imag.array = out_vec._data[:, i].imag
                    self._scatter(out_petsc_imag, in_petsc_imag, addv=False, mode=False)

                    in_vec._data[:, i] = in_petsc.array + in_petsc_imag.array * 1j

            else:
                for i in range(in_vec._ncol):
                    in_petsc.array = in_vec._data[:, i]
                    out_petsc.array = out_vec._data[:, i]
                    self._scatter(out_petsc, in_petsc, addv=False, mode=False)
                    in_vec._data[:, i] = in_petsc.array

        elif mode == 'rev':
            in_petsc = in_vec._petsc
            out_petsc = out_vec._petsc
            for i in range(in_vec._ncol):
                in_petsc.array = in_vec._data[:, i]
                out_petsc.array = out_vec._data[:, i]
                self._scatter(in_petsc, out_petsc, addv=True, mode=True)
                out_vec._data[:, i] = out_petsc.array
