# -*- coding: utf-8 -*-
###########################################################################
# Copyright (c), The AiiDA team. All rights reserved.                     #
# This file is part of the AiiDA code.                                    #
#                                                                         #
# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
# For further information on the license, see the LICENSE.txt file        #
# For further information please visit http://www.aiida.net               #
###########################################################################
"""Data plugin to represet arrays of projected wavefunction components."""
import copy

import numpy as np

from aiida.common import exceptions
from aiida.plugins import OrbitalFactory

from ..orbital import OrbitalData
from .array import ArrayData
from .bands import BandsData

__all__ = ('ProjectionData',)


class ProjectionData(OrbitalData, ArrayData):
    """
    A class to handle arrays of projected wavefunction data. That is projections
    of a orbitals, usually an atomic-hydrogen orbital, onto a
    given bloch wavefunction, the bloch wavefunction being indexed by
    s, n, and k. E.g. the elements are the projections described as
    < orbital | Bloch wavefunction (s,n,k) >
    """

    def _check_projections_bands(self, projection_array):
        """
        Checks to make sure that a reference bandsdata is already set, and that
        projection_array is of the same shape of the bands data

        :param projwfc_arrays: nk x nb x nwfc array, to be
                               checked against bands

        :raise: AttributeError if energy is not already set
        :raise: AttributeError if input_array is not of same shape as
                dos_energy
        """
        try:
            shape_bands = np.shape(self.get_reference_bandsdata())
        except AttributeError:
            raise exceptions.ValidationError('Bands must be set first, then projwfc')
        # The [0:2] is so that each array, and not collection of arrays
        # is used to make the comparison
        if np.shape(projection_array) != shape_bands:
            raise AttributeError('These arrays are not the same shape as the bands')

    def set_reference_bandsdata(self, value):
        """
        Sets a reference bandsdata, creates a uuid link between this data
        object and a bandsdata object, must be set before any projection arrays

        :param value: a BandsData instance, a uuid or a pk
        :raise: exceptions.NotExistent if there was no BandsData associated with uuid or pk
        """
        from aiida.orm import load_node

        if isinstance(value, BandsData):
            uuid = value.uuid
        else:
            try:
                pk = int(value)
                bands = load_node(pk=pk)
                uuid = bands.uuid
            except ValueError:
                uuid = str(value)
                try:
                    bands = load_node(uuid=uuid)
                    uuid = bands.uuid
                except Exception:  # pylint: disable=bare-except
                    raise exceptions.NotExistent(
                        'The value passed to set_reference_bandsdata was not associated to any bandsdata'
                    )

        self.set_attribute('reference_bandsdata_uuid', uuid)

    def get_reference_bandsdata(self):
        """
        Returns the reference BandsData, using the set uuid via
        set_reference_bandsdata

        :return: a BandsData instance
        :raise AttributeError: if the bandsdata has not been set yet
        :raise exceptions.NotExistent: if the bandsdata uuid did not retrieve bandsdata
        """
        from aiida.orm import load_node
        try:
            uuid = self.get_attribute('reference_bandsdata_uuid')
        except AttributeError:
            raise AttributeError('BandsData has not been set for this instance')
        try:
            bands = load_node(uuid=uuid)
        except exceptions.NotExistent:
            raise exceptions.NotExistent('The bands referenced to this class have not been found in this database.')
        return bands

    def _find_orbitals_and_indices(self, **kwargs):
        """
        Finds all the orbitals and their indicies associated with kwargs
        essential for retrieving the other indexed array parameters

        :param kwargs: kwargs that can call orbitals as in get_orbitals()
        :return: retrieve_indexes, list of indicicies of orbitals corresponding
                 to the kwargs
        :return: all_orbitals, list of orbitals to which the indexes correspond
        """
        selected_orbitals = self.get_orbitals(**kwargs)
        selected_orb_dicts = [orb.get_orbital_dict() for orb in selected_orbitals]
        all_orbitals = self.get_orbitals()
        all_orb_dicts = [orb.get_orbital_dict() for orb in all_orbitals]
        retrieve_indices = [i for i in range(len(all_orb_dicts)) if all_orb_dicts[i] in selected_orb_dicts]
        return retrieve_indices, all_orbitals

    def get_pdos(self, **kwargs):
        """
        Retrieves all the pdos arrays corresponding to the input kwargs

        :param kwargs: inputs describing the orbitals associated with the pdos
                       arrays
        :return: a list of tuples containing the orbital, energy array and pdos
                 array associated with all orbitals that correspond to kwargs

        """
        retrieve_indices, all_orbitals = self._find_orbitals_and_indices(**kwargs)
        out_list = [(
            all_orbitals[i], self.get_array(f'pdos_{self._from_index_to_arrayname(i)}'),
            self.get_array(f'energy_{self._from_index_to_arrayname(i)}')
        ) for i in retrieve_indices]
        return out_list

    def get_projections(self, **kwargs):
        """
        Retrieves all the pdos arrays corresponding to the input kwargs

        :param kwargs: inputs describing the orbitals associated with the pdos
                       arrays
        :return: a list of tuples containing the orbital, and projection arrays
                 associated with all orbitals that correspond to kwargs

        """
        retrieve_indices, all_orbitals = self._find_orbitals_and_indices(**kwargs)
        out_list = [
            (all_orbitals[i], self.get_array(f'proj_{self._from_index_to_arrayname(i)}')) for i in retrieve_indices
        ]
        return out_list

    @staticmethod
    def _from_index_to_arrayname(index):
        """
        Used internally to determine the array names.
        """
        return f'array_{index}'

    def set_projectiondata(
        self,
        list_of_orbitals,
        list_of_projections=None,
        list_of_energy=None,
        list_of_pdos=None,
        tags=None,
        bands_check=True
    ):
        """
        Stores the projwfc_array using the projwfc_label, after validating both.

        :param list_of_orbitals: list of orbitals, of class orbital data.
                                 They should be the ones up on which the
                                 projection array corresponds with.

        :param list_of_projections: list of arrays of projections of a atomic
                              wavefunctions onto bloch wavefunctions. Since the
                              projection is for every bloch wavefunction which
                              can be specified by its spin (if used), band, and
                              kpoint the dimensions must be
                              nspin x nbands x nkpoints for the projwfc array.
                              Or nbands x nkpoints if spin is not used.

        :param energy_axis: list of energy axis for the list_of_pdos

        :param list_of_pdos: a list of projected density of states for the
                             atomic wavefunctions, units in states/eV

        :param tags: A list of tags, not supported currently.

        :param bands_check: if false, skips checks of whether the bands has
                            been already set, and whether the sizes match. For
                            use in parsers, where the BandsData has not yet
                            been stored and therefore get_reference_bandsdata
                            cannot be called
        """

        # pylint: disable=too-many-arguments,too-many-locals,too-many-branches,too-many-statements

        def single_to_list(item):
            """
            Checks if the item is a list or tuple, and converts it to a list
            if it is not already a list or tuple

            :param item: an object which may or may not be a list or tuple
            :return: item_list: the input item unchanged if list or tuple and
                                [item] otherwise
            """
            if isinstance(item, (list, tuple)):
                return item

            return [item]

        def array_list_checker(array_list, array_name, orb_length):
            """
            Does basic checks over everything in the array_list. Makes sure that
            all the arrays are np.ndarray floats, that the length is same as
            required_length, raises exception using array_name if there is
            a failure
            """
            if not all(isinstance(_, np.ndarray) for _ in array_list):
                raise exceptions.ValidationError(f'{array_name} was not composed entirely of ndarrays')
            if len(array_list) != orb_length:
                raise exceptions.ValidationError(f'{array_name} did not have the same length as the list of orbitals')

        ##############
        list_of_orbitals = single_to_list(list_of_orbitals)
        list_of_orbitals = copy.deepcopy(list_of_orbitals)

        # validates the input data
        if not list_of_pdos and not list_of_projections:
            raise exceptions.ValidationError('Must set either pdos or projections')
        if bool(list_of_energy) != bool(list_of_pdos):
            raise exceptions.ValidationError('list_of_pdos and list_of_energy must always be set together')

        orb_length = len(list_of_orbitals)

        # verifies and sets the orbital dicts
        list_of_orbital_dicts = []
        for i, _ in enumerate(list_of_orbitals):
            this_orbital = list_of_orbitals[i]
            orbital_dict = this_orbital.get_orbital_dict()
            try:
                orbital_type = orbital_dict.pop('_orbital_type')
            except KeyError:
                raise exceptions.ValidationError(f'No _orbital_type key found in dictionary: {orbital_dict}')
            cls = OrbitalFactory(orbital_type)
            test_orbital = cls(**orbital_dict)
            list_of_orbital_dicts.append(test_orbital.get_orbital_dict())
        self.set_attribute('orbital_dicts', list_of_orbital_dicts)

        # verifies and sets the projections
        if list_of_projections:
            list_of_projections = single_to_list(list_of_projections)
            array_list_checker(list_of_projections, 'projections', orb_length)
            for i, _ in enumerate(list_of_projections):
                this_projection = list_of_projections[i]
                array_name = self._from_index_to_arrayname(i)
                if bands_check:
                    self._check_projections_bands(this_projection)
                self.set_array(f'proj_{array_name}', this_projection)

        # verifies and sets both pdos and energy
        if list_of_pdos:
            list_of_pdos = single_to_list(list_of_pdos)
            list_of_energy = single_to_list(list_of_energy)
            array_list_checker(list_of_pdos, 'pdos', orb_length)
            array_list_checker(list_of_energy, 'energy', orb_length)
            for i, _ in enumerate(list_of_pdos):
                this_pdos = list_of_pdos[i]
                this_energy = list_of_energy[i]
                array_name = self._from_index_to_arrayname(i)
                if bands_check:
                    self._check_projections_bands(this_projection)
                self.set_array(f'pdos_{array_name}', this_pdos)
                self.set_array(f'energy_{array_name}', this_energy)

        # verifies and sets the tags
        if tags is not None:
            try:
                if len(tags) != len(list_of_orbitals):
                    raise exceptions.ValidationError('must set as many tags as projections')
            except IndexError:
                return exceptions.ValidationError('tags must be a list')

            if not all(isinstance(_, str) for _ in tags):
                raise exceptions.ValidationError('Tags must set a list of strings')
            self.set_attribute('tags', tags)

    def set_orbitals(self, **kwargs):  # pylint: disable=arguments-differ
        """
        This method is inherited from OrbitalData, but is blocked here.
        If used will raise a NotImplementedError
        """
        raise NotImplementedError(
            'You cannot set orbitals using this class!'
            ' This class is for setting orbitals and '
            ' projections only!'
        )
