"""
======================================
  ViPErLEED Graphical User Interface
======================================
 *** module guilib.leedsim.classes.oldleedpattern ***

Defines the LEEDPattern and LEEDsubpattern classes, used for displaying a
LEEDPattern. This is the old version of the classes. Reimplemented with new
logics in leedpattern.py. This module is kept until the reimplementation is
ready.

Author: Michele Riva
Created: 2021-03-13
"""

import numpy as np
from quicktions import Fraction
from scipy import spatial as sp
from matplotlib import cm  # color maps
from matplotlib import colors as mpl_colors

from viperleed import guilib as gl


class LEEDPattern:
    # @gl.profile_calls(print_args=[30])
    def __init__(self, leed_parameters):
        # run leed_parameters through gl.LEEDParameters, so that all
        # parameters are correctly interpreted, and missing optional values
        # are set to their default
        self.__parameters = gl.LEEDParameters(leed_parameters)

        # The following attributes are set by __build_lattices() and
        # __build_LEEDPattern()
        # * reciprocal-space lattices for surface and bulk
        self.reciprocal_lattices = {'surf': None, 'bulk': None}

        # * characteristics of domains generated by symmetry operations
        #   of the bulk. 'g' is a list (one element per domain) containing
        #   all the reciprocal-space wave-vectors (in x,y Cartesian coordinates)
        #   of the beams from one domain
        self.domains = {'operations': None, 'g': None}

        ## NEED TO REDUCE THE NUMBER OF ATTRIBUTES FROM HERE ON
        self.surfLEED = None  # like self.domains['g'], but flattened
        self.superposedBeams = None  # grouping of beam indices telling which ones overlap   ## PROBABLY WILL GO with BeamIndexList
        self.names = None     # strings of beam names (wrt bulk basis)            ## REPLACE via BeamIndexList!

        # The following one is related to symmetry-equivalent beams
        self.firstDomRefs = None  # references of equivalent beams in 1st domain  ## REPLACE (different logics for off normal).
        
        # The following ones are related to the LEED subpatterns                  ## THE LOGICS NEEDS TO CHANGE!
        # Will probably have a self.patterns dictionary that is initially populated only with some of the subpatterns. Other combinations will be dynamically generated and stored as new keys, with some reasonable hashing.
        self.firstLEED = None  # subpattern of first domain
        self.domsLEED = None   # subpatterns for all domains
        self.domColors = None  # colors for plotting multiple domains

        self.__build_lattices()
        self.__build_LEEDPattern()

    @property
    def bulk_basis(self):
        if self.reciprocal_lattices['bulk'] is None:
            return np.dot(self.superlattices[0].T,                              # ENOUGH TO HAVE THE FIRST ONE
                          self.reciprocal_lattices['surf'].basis)
        return self.reciprocal_lattices['bulk'].basis

    @property
    def max_energy(self):
        """
        Maximum LEED energy
        """
        return self.__parameters['eMax']

    @property
    def n_beams(self):                                                          # NOT CORRECT FOR MULTIPLE STRUCTURES. Probably not needed
        """
        Number of LEED beams in the pattern
        """
        return self.reciprocal_lattices['surf'].n_beams

    @property
    def n_domains(self):
        """
        Number of distinct domains produced by the symmetry operations of the
        bulk
        """
        return len(self.domains['operations'])

    @property
    def superlattices(self):
        """
        Numpy ndarray of superlattice matrices generating the symmetry-related
        domains. The first element is the one generating the lattice whose basis
        is given in the constructor as a LEED parameter
        """
        superlattice = self.__parameters['SUPERLATTICE']
        operations = self.domains['operations']
        if operations is None:
            operations = [gl.PlaneGroup.E]  # identity matrix
        # return np.asarray([np.dot(superlattice, operation)
                           # for operation in operations])
        return np.einsum('ij,mjk->mik', superlattice, operations)
    
    @property
    def primary_beam_theta(self):
        """
        Polar angle of incidence with respect to the direction perpendicular
        to the surface
        """
        return self.__parameters['beamIncidence'][0]
    
    @property
    def primary_beam_phi(self):
        """
        Azimuthal angle of incidence of the primary beam with respect to the
        x axis. Positive counterclockwise when looking down on the surface.
        """
        return self.__parameters['beamIncidence'][1]
    
    @property
    def primary_beam_direction(self):                                           # will probably not even use this
        """
        2-tuple representing the in-plane direction vector of incidence,
        in Cartesian coordinates
        """
        if self.primary_beam_theta == 0:
            return (0, 0)
        phi = np.radians(self.primary_beam_phi)
        return np.cos(phi), np.sin(phi)

    def __build_lattices(self):                                                 # Not even sure if we really need both lattices anymore. Probably bulk + domains are enough
        # Create dummy surface lattice just to get the reciprocal basis
        surf = gl.Lattice(self.__parameters['surfBasis'])

        # and get the maximum screen radius
        max_radius = self.screen_radius(self.max_energy)

        # Now make the reciprocal lattices:
        self.reciprocal_lattices['surf'] = gl.Lattice(
            surf.reciprocal_basis, space='reciprocal',
            group=self.__parameters['surfGroup'], limit=max_radius
            )
        self.reciprocal_lattices['bulk'] = gl.Lattice(
            self.bulk_basis, space='reciprocal',
            group=self.__parameters['bulkGroup'], limit=max_radius
            )
        self.reciprocal_lattices['bulk'].group.screws_glides = (
            self.__parameters['bulk3Dsym'],
            self.reciprocal_lattices['bulk'].cell_shape
            )

    # @gl.profile_lines
    def __build_LEEDPattern(self):                                                ## bottleneck!

        # 1) get the (bulk) operations that generate symmetry-equivalent domains
        self.domains['operations'] = self.__find_domain_operations()

        # 2) generate the gx gy coordinates for all the domains                   ## PERHAPS IT MAKES SENSE TO GENERATE hk(s) FIRST? THESE CAN BE CONVERTED TO a BeamIndexList
        self.domains['g'] = [                                                     ## NEEDS TO CHANGE TO ACCOUNT FOR OFF-NORMAL. It may also be that I don't need this at all,
            np.linalg.multi_dot(
                (self.reciprocal_lattices['surf'].hk,
                 np.linalg.inv(superlattice).T,
                 self.bulk_basis)
                )
            for superlattice in self.superlattices                                ## CAN BE PART OF LEEDDomains
            ]

        # 3) prepare the full LEED pattern of the surface, including all
        # domains
        surfLEEDx = np.array([self.domains['g'][i][:, 0]
                              for i in range(self.n_domains)]).flatten()
        surfLEEDy = np.array([self.domains['g'][i][:, 1]
                              for i in range(self.n_domains)]).flatten()
        self.surfLEED = np.array(list(zip(surfLEEDx, surfLEEDy)))

        # 4) find superposed beams
        self.superposedBeams = self.get_supSpots()

        # 5) prepare the formatted names of the fractional indices
        self.names = np.array([
            self.formatFractionalIndices(self.reciprocal_lattices['surf'].hk, m)  # MAIN BOTTLENECK: > half time of whole method
            for m in self.superlattices])                                         # WILL NOT EXIST

        # 6) and build the subpatterns
        self.__build_subpatterns()                                                ## BOTTLENECK: half of time of whole method

    def get_supSpots(self):                                                       ## TO EVALUATE WHETHER IT'S FASTER TO USE BeamIndexList
        # Group all the beams into separate groups, such that each group
        # contains the indices of overlapping beams (can be one index if non
        # overlapping)
        tree = sp.cKDTree(self.surfLEED)
        return np.array(tree.query_ball_point(x=self.surfLEED, r=1e-8))

    # @gl.profile_calls(print_args=[30])
    def get_beamGrouping(self, domains=None):                                     ## TO RETHINK FOR NON-NORMAL INCIDENCE  ## BOTTLENECK for __build_subpatterns
        # this is wrong for non-normal incidence!
        # all domains should be treated separately, rather than using the
        # symmetry of the first one, since the special direction given by the
        # beam incidence has a different effect on each domain
        self.firstDomRefs = self.get_FirstDomainSymmetry()                        ## BOTTLENECK?

        if domains is None:
            domains = range(self.n_domains)

        # beamsRefs contains at index [i, j] a list of the indices of beams of
        # domain domains[i] equivalent to beam j (including j)
        beamsRefs = []
        for dom in domains: 
            beamsRefs.append([crossRef + dom*self.n_beams
                              for crossRef in self.firstDomRefs])
        return np.array(beamsRefs, dtype=object)

    def get_equivalentSpots(self, domains=None):                                  ## REPLACE with generic function!
        if not hasattr(domains, '__len__') and domains is not None:
            raise TypeError("Invalid type for domains. "
                            "Should be None or an array-like.")
        if domains is None:
            # output spots from all domains
            domains = range(self.n_domains)
        elif not all(isinstance(dom, int) for dom in domains):
            raise ValueError("Invalid domain index. "
                             "All indices should be integers")
        elif any(dom < 0 or dom >= self.n_domains for dom in domains):
            raise ValueError("Domain index out of range. "
                             "Indices should be between 0 "
                             f"and {self.n_domains}.")

        names = self.names[domains]
        beamRefs = self.get_beamGrouping(domains)

        # get ready to prepare one dictionary per each domain:
        # key: beam name
        # value: names of beams equivalent to key
        #        Notice that for each key (=beam) in one domain, one should
        #        also include as equivalent beams those from other domains
        #        at the same positional index.
        #        !!!THIS IS A PROBLEM FOR NON-NORMAL INCIDENCE!!!
        #        --> get this list (the same for all domains) ready first
        allRefs = [np.concatenate(beamRefs[:, beam]).astype(int)
                   for beam in range(self.n_beams)]
        refNames = [set(self.names.ravel()[beamRef]) for beamRef in allRefs]

        # now pack the dictionaries
        # Notice that it is necessary to have one dictionary per domain since
        #   some domains will give rise to a subset of all the LEED spots:
        #   e.g., in a (2x1) on square bulk, one domain will give spot (1/2, 0)
        #   the other domain spot (0, 1/2)
        # Having all the dictionaries allows to determine which of the
        #   superposed spots are indeed equivalent (see next comment) as well
        #   as which domains overlap
        domsDicts = [dict(zip(domNames, refNames)) for domNames in names]

        # now figure out which of the beams are symmetry equivalent
        # taking superposition of spots into account:
        # spots originating from beams that are symmetry equivalent within
        # each domain, are symmetry equivalent only if all domains contribute
        # with symmetry-equivalent beams in all spots.
        #
        # For example (hexagonal bulk):
        # dom1: (1 0), (-1 0), (1 -1), (-1 1) equivalent
        # dom2: (1 0), (-1 0), (0 -1), (0 1)  equivalent
        # dom3: (1 0), (-1 0)                 equivalent
        # --> (1 0) and (-1 0) are equivalent
        #
        # dom1: (1 0),  (-1 0),  (1 -1), (-1 1)    equivalent
        # dom2:                  (1 -1), (-1 1)    equivalent
        # dom3: (0, 1), (0, -1), (1 -1), (-1 1)    equivalent
        # --> (1 -1) and (-1 1) are equivalent, but are
        #     not equivalent to [(1 0), (-1 0)]
        eqBeams = []
        for name in names.ravel():
            setLst = [dDict[name] for dDict in domsDicts if name in dDict]
            if setLst:
                newBeams = set.intersection(*setLst)
                # remove the keys of the beams that have
                # already been processed from each domain
                for domDict in domsDicts:
                    for beam in newBeams:
                        if beam in domDict:
                            del domDict[beam]
                eqBeams.append(newBeams)

        # sort within each equivalence group
        lstBeams = [sorted(list(beams), key=self.beamsSortCriterion)
                    for beams in eqBeams]
        # and by energy
        sortedBeams = sorted(lstBeams, key=self.sortEnergy)

        beamsWithIndices = []
        overlapDoms = []
        for (i, beams) in enumerate(sortedBeams):
            for beam in beams:
                overlapDoms = [dom + 1 for dom in domains
                               if beam in self.names[dom].ravel()]
                beamsWithIndices.append([beam, i, overlapDoms])

        # Now find which of the beams are extinct due to glide, place them
        # in a list of beam indices in the same order as in self.names[i]
        extinct = self.get_ExtinctFirstDomain()

        extinctNames = self.names[:, extinct]
        # notice that extinctNames is an empty list if there are no glide
        # extinct beams.
        for b, (beam, group, doms) in enumerate(beamsWithIndices):
            extDoms = [dom for dom in doms if beam in extinctNames[dom-1]]
            # Notice the dom-1 because of the way the doms list is
            # built in the previous for loop

            beamsWithIndices[b] = [beam, group, doms, extDoms]
            # now, only in case extDoms is of length 1, i.e., one domain
            # only gives rise to the extinct beam, the group index goes
            # negative
            # !!! Perhaps will be changed with the better version in base.py !!!
            if len(extDoms) == 1:
                beamsWithIndices[b][1] *= -1

        return beamsWithIndices

    def beamsSortCriterion(self, beamstr):                                        ## this may go into BeamIndexList
        beam = eval(beamstr)  # returns a tuple                                   ## This will go

        sortH = beam[0]
        sortK = beam[1]

        return (-sortH-sortK, -sortH)
    
    def sortEnergy(self, beamlist):                                               ## this may go into BeamIndexList
        beam = eval(beamlist[0])  # returns a tuple
        # sort by Energy, i.e., by g^2. Notice the
        # rounding to reduce floating-point errors
        g_norm = np.linalg.norm(np.dot(beam, self.bulk_basis)).round(7)
        return (g_norm, *self.beamsSortCriterion(beamlist[0]))

    # @gl.profile_lines
    def __build_subpatterns(self):
        # Now sort the sub-patterns:
        # - There will always be a firstLEED and a domsLEED
        #   This is to handle the plotting of the first domain only.              ## NEEDS TO CHANGE TO DISPLAY EACH DOMAIN SEPARATELY OR AN ARBITRARY COMBINATION OF DOMAINS
        #   All entries in these lists are of type LEEDsubpattern
        #
        # - firstLEED, contains at most two entries: the first one
        #   includes all non-extinct beams, the second one
        #   the extinct ones, if any are present
        # - domsLEED is empty if only one domain is present.
        #   Otherwise, it contains:
        #   * one entry with superimposed spots
        #   * at most two entries for each domain.
        #     first entry: non-extinct; second entry: extinct (if any)

        doms = self.domains['g']
        symEq = self.get_beamGrouping()                                          # BOTTLENECK
        names = self.names

        maskFirst = np.full(self.n_beams, True)
        maskFirst[self.get_ExtinctFirstDomain()] = False

        self.firstLEED = [LEEDsubpattern(doms[0][maskFirst],
                                         symEq[0][maskFirst],
                                         names[0][maskFirst])]
        if all(maskFirst):  # no glide-extinct beams
            pass
        else:
            glideAlpha = 0.2
            glideScale = 2
            maskExtinct = np.invert(maskFirst)
            self.firstLEED.append(LEEDsubpattern(doms[0][maskExtinct],
                                                 symEq[0][maskExtinct],
                                                 names[0][maskExtinct],
                                                 color='k',
                                                 alpha=glideAlpha,
                                                 marker='x',
                                                 sizeScale=glideScale
                                                 ))

        self.domsLEED = []
        self.domColors = None
        if self.n_domains > 1:
            # If there are domains, there always will be overlapping spots,
            # at least those at integer orders

            # superposedMask is True at all the beam indices of beams that
            # are superposed to others
            superposedMask = np.array([len(sup) > 1
                                       for sup in self.superposedBeams])

            # figure out which domains superpose:
            overlapDoms = [[x//self.n_beams + 1 for x in sup]
                           for sup in self.superposedBeams[superposedMask]]
            self.domsLEED = [LEEDsubpattern(self.surfLEED[superposedMask],
                                            symEq.ravel()[superposedMask],
                                            names.ravel()[superposedMask],
                                            domain=-1,
                                            color='gray',
                                            overlapDoms=overlapDoms
                                            )]

            # define colors
            self.domColors = cm.gnuplot(np.linspace(0.1, 0.9, self.n_domains))

            for dom in range(self.n_domains):
                mask = np.invert(
                    superposedMask[dom*self.n_beams:(dom+1)*self.n_beams])
                maskDom = np.logical_and(maskFirst, mask)
                self.domsLEED.append(LEEDsubpattern(doms[dom][maskDom],
                                                    symEq[dom][maskDom],
                                                    names[dom][maskDom],
                                                    domain=dom+1,
                                                    color=self.domColors[dom]
                                                    )
                                     )
                if all(maskFirst):  # no glide extinct
                    pass
                else:
                    maskDomExtinct = np.logical_and(maskExtinct, mask)
                    self.domsLEED.append(
                        LEEDsubpattern(doms[dom][maskDomExtinct],
                                       symEq[dom][maskDomExtinct],
                                       names[dom][maskDomExtinct],
                                       domain=dom+1,
                                       color=self.domColors[dom],
                                       alpha=glideAlpha,
                                       marker='x',
                                       sizeScale=glideScale))

    def screen_radius(self, en):
        return gl.screen_radius(en, self.__parameters['screenAperture'])

    def get_FirstDomainSymmetry(self):
        """
        Returns
        -------
        crossRefs: list of np.arrays,
                   length == len(self.reciprocal_lattices['surf'].hk)
                   Each np.array contains the positional indices of the
                   beams equivalent to the current one (including
                   self-reference)
        """

        hk = self.reciprocal_lattices['surf'].hk
        N = len(hk)
        symOps = self.reciprocal_lattices['surf'].group.operations()

        # at first, use sets to avoid duplicates
        crossRefs = [{i} for i in range(N)]

        for op in symOps:
            hk_transformed = np.dot(op, hk.T).T.tolist()
            # transform hk according to the operation of the group
            for i in range(N):
                idx = hk_transformed.index(hk[i].tolist())
                # find which transformed hk is equal to the current beam
                # NB: there's always only one for each (beam, operation) pair
                #
                # and add it to the set of cross references for the current
                # beam
                crossRefs[i] |= {idx}

        # then convert each set to a np.array of int, so that it's iterable
        # and it's easy to process later
        return np.array([np.array(list(x), dtype=int) for x in crossRefs], 
                        dtype=object)

    def get_ExtinctFirstDomain(self):
        """
        Returns
        -------
        list of indices of beams of the first domain that are extinct
        due to glide
        """
        hk = self.reciprocal_lattices['surf'].hk
        group = self.reciprocal_lattices['surf'].group.group

        if 'g' in group:  # group has glide
            if '[1 0]' in group:  # pg or pmg
                idx = np.where([x[1] == 0 and abs(x[0] % 2) == 1
                                for x in hk])[0].tolist()
            elif '[0 1]' in group:  # pg or pmg
                idx = np.where([x[0] == 0 and abs(x[1] % 2) == 1
                                for x in hk])[0].tolist()
            else:  # pgg or p4g
                idx = np.where([(x[1] == 0 and abs(x[0] % 2) == 1)
                                or (x[0] == 0 and abs(x[1] % 2) == 1)
                                for x in hk])[0].tolist()
        else:
            idx = []

        return idx

    def __find_domain_operations(self):
        """
        Finds the symmetry operations of the bulk lattice that give distinct
        domains in LEED.
        NB: this function gives different results than what one would get from
        LEEDPat, as we're interested also in the symmetry relations between the
        intensities of the spots, while LEEDPat cares only about the presence
        or not of any spot.

        Returns
        -------
        list of operations

        -------
        """
        # The current version is based on the concept of co-sets of a group.
        # Given a group G and a subgroup H, the left co-set of H with respect
        # to the group operation g of G is
        #           gH = {g*h : h in H}.
        # Additionally, we take into account that each element of G is found in
        # exactly only one co-set (e.g., H is the identity co-set).
        #
        # The operations g_i that generate distinct co-sets are those that will
        # generate distinct lattices

        def __array_2_tuple(arr):
            """
            Convenience function that is used to convert a 2x2 array to a 2x2
            tuple
            """
            return tuple(map(tuple, arr))

        # 1) project the operations from the surface group to the bulk; also
        #    rounding to integers
        transform = np.linalg.inv(self.superlattices[0])                        # FIRST ONE IS ENOUGH
        surf_ops = tuple(
            op.round().astype(int)
            for op
            in self.reciprocal_lattices['surf'].group.transform(transform)
            )

        # 2) keep track of the operations of the bulk group that are already in
        #    a co-set, and of the group operations giving distinct co-sets
        coset_ops = set()
        _ops = []

        # 3) run through the bulk operations, adding the operations of new
        #    co-sets to the set above
        bulk_group = self.reciprocal_lattices['bulk'].group
        bulk_ops = bulk_group.operations(include_3d=True)

        for b_op in bulk_ops:
            if b_op in coset_ops:
                continue
            _ops.append(b_op)
            coset = set(__array_2_tuple(np.dot(b_op, surf_op))
                        for surf_op in surf_ops)
            coset_ops.update(coset)

        return _ops

    def formatFractionalIndices(self, hk, m):                                     ## this will probably go, and will rather use BeamIndex; BOTTLENECK
        """
        Method that generates fractional indices of a superstructure m whose
        integer indices are hk.

        Parameters
        ----------
        hk, 2D-array
            Integer indices along the two directions to be formatted

        m, 2x2 np.array
            superlattice matrix

        Returns
        -------
        names, 1D array
            each element is a string with formatted names for the indices
        """

        # The beam indices [hb,kb] for the surface indices [h,k] are
        #    [hb, kb] = np.dot([h, k], G),
        # with G = np.linalg.inv(M).transpose()
        hkbulk = np.dot(np.linalg.inv(m), hk.T).T

        # now get the formatting done:
        # all fractional indices are of type nn/mu with nn integer, and
        mu = abs(np.round(np.linalg.det(m)).astype(int)).item()

        # notice that, to prevent Fraction().limit_denominator() from
        # overflowing I'm casting the numpy-type numbers to their native
        # python type via .item()
        names = [', '.join([str(Fraction(hh.item()).limit_denominator(mu))         ## BOTTLENECK is limit_denominator
                            for hh in hkb]) for hkb in hkbulk]

        return np.array(names)

    def rotate(self, angle, which='surf'):
        """
        This function is a mess and needs to be split into two functions.
        It currently does two different things: if which='bulk', it returns a
        copy of the bulk lattice. If which='surf' it rotates all the LEED
        subpatterns. It would probably be nicer to also have the bulk pattern
        be a special subpattern.
        """
        which = which[0:4]
        if which not in ['surf', 'bulk']:
            raise ValueError("LEEDPattern: Only 'surf', 'surface', or "
                             "'bulk' are acceptable values for rotate()")
        if hasattr(angle, '__len__'):
            raise ValueError('Angle input is not a scalar')

        if which == 'surf':
            angle = np.radians(angle)
            rot = np.array([[np.cos(angle), np.sin(angle)],
                            [-np.sin(angle), np.cos(angle)]])
            for pat in [*self.firstLEED, *self.domsLEED]:
                pat.transformCoords(rot)
            return None
        return self.reciprocal_lattices['bulk'].get_rotated_lattice(angle)


class LEEDsubpattern:
    """
    --> I originally intended to use this guy for the
         determination of beam equivalence and
         handling of hovering annotations
    """
    def __init__(self, beams, symBeams, names, domain=1, color='k',
                 alpha=None, marker='o', sizeScale=1, overlapDoms=None):
        if not all([len(beams) == len(x) for x in [symBeams, names]]):
            raise ValueError('Input shape mismatch. The length of symBeams'
                             'and/or _names differs from the one of beams.')
        if not isinstance(domain, int):
            raise ValueError("domain must be an integer.")

        self.beams = beams
        self.rotBeams = self.beams
        self.symmetricBeams = symBeams
        self.names = names
        self.domain = domain
        if alpha is None:
            self.color = np.array([mpl_colors.to_rgba(color, alpha=1)])
        else:
            self.color = np.array([mpl_colors.to_rgba(color, alpha=alpha)])
        self.marker = marker
        self.sizeScale = sizeScale
        self.overlappingDomains = overlapDoms

    def transformCoords(self, transf):
        if np.shape(transf) == (2, 2):
            self.rotBeams = np.dot(self.beams, transf)
        else:
            raise

