"""
@package pmsco.transforms.multipoles
Transform holoscan data to multipoles expansion.

@author Matthias Muntwiler

@copyright (c) 2015-24 by Paul Scherrer Institut @n
Licensed under the Apache License, Version 2.0 (the "License"); @n
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  http://www.apache.org/licenses/LICENSE-2.0
"""

import logging
import numpy as np
import numpy.typing as npt
import os
import scipy.special
from typing import Any, Callable, Dict, Generator, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, Union
import h5py

from pmsco.data import analyse_holoscan_steps, create_data, DATATYPES, DTYPE_TPI

logger = logging.getLogger(__name__)


class MultipoleExpansion:
    """
    Multipole expansion (spherical harmonics) of hologram scans

    This class provides functions to calculate the multipole expansion of hologram scans.

    Usage
    -----

    1. Instantiate the object
    2. Assign a scan to `holoscan`.
       If you will use the `generate` method, the scan must include an intensity column.
       Else, the theta and phi coordinates are sufficient.
    3. Set `lmax`.
    4. Call `generate` to calculate the alm coefficients, or assign the coefficients directly.
    5. Call `expand` to calculate the filtered pattern.

    The object keeps intermediate calculation data, e.g. spherical harmonics, to speed up further calculations.
    You have to keep the object alive in order to benefit from these performance improvements.
    Reuse of the spherical harmonics depends on the same theta-phi scan grid and same lmax.
    The class re-calculates the internal data automatically if these conditions are not met.

    Attributes
    ----------

    _sph_harm dimensions: l, m, tp
    _d_omege dimensions: l, m, tp
    _alm dimensions: l, m

    """

    def __init__(self) -> None:
        """
        Initialize the object

        The scan and arrays are set to shape 0.
        The default lmax is 30 and the alm array is initialized with zeros to the corresponding shape.
        The spherical harmonics and volume element are not initialized.
        """

        super().__init__()
        self._scan: np.ndarray = create_data((0,), datatype='TP')
        self._original: np.ndarray = np.zeros(self._scan.shape)
        self._expansion: np.ndarray = np.zeros(self._scan.shape)
        self._lmax: int = 60
        self._alm: np.ndarray = np.zeros((self._lmax // 2 + 1, 2 * self._lmax + 1))
        self._sph_harm: Optional[np.ndarray] = None
        self._d_omega: Optional[np.ndarray] = None

    @property
    def holoscan(self) -> np.ndarray:
        """
        Holoscan grid and original intensity pattern.

        The holoscan array can be created, e.g., from holo_grid or is taken from measured data.
        't' and 'p' contain the polar and azimuthal coordinates of the scan.
        The 'i' column contains the intensity value.

        The array must contain at least the 't' and 'p' columns if you want to expand the pattern from alm.
        If you want to generate the coefficients, the 'i' (intensity) column must be set as well.

        @attention: The property returns a copy of internal data.
        Do not change individual values or columns directly on the array.
        """

        tpi = np.zeros(self._scan.shape, dtype=DTYPE_TPI)
        tpi['t'] = self._scan['t']
        tpi['p'] = self._scan['p']
        tpi['i'] = self._original

        return tpi

    @holoscan.setter
    def holoscan(self, arr: np.ndarray):
        if arr.shape != self._scan.shape or \
                self._original.shape[0] == 0 or \
                np.all(np.abs(arr['i'] - self._original) < 0.01):
            self._scan = create_data(arr.shape, datatype='TP')
            self._discard_basis()

        self._original = arr['i']
        self._scan['t'] = arr['t']
        self._scan['p'] = arr['p']

    @property
    def expansion(self) -> np.ndarray:
        """
        Return the expansion in TPI format.

        The expansion must have been calculated using the `expand` method.

        @return: Structured numpy array in TPI format (dtype=DTYPE_TPI).
        """

        tpi = create_data(self._expansion.shape, dtype=DTYPE_TPI)
        tpi['t'] = self._scan['t']
        tpi['p'] = self._scan['p']
        tpi['i'] = self._expansion

        return tpi

    def get_expansion(self, datatype: Optional[str] = None, dtype: Optional[npt.DTypeLike] = None) -> np.ndarray:
        """
        Return the expansion in the requested format.

        The expansion must have been calculated using the `expand` method.
        Either datatype or dtype must be specified, dtype takes precedence.

        @param datatype       see DATATYPES.
        @param dtype          see DTYPES.
                              Any structured array type will work.
                              The dtype must specify at least the `i`, `t` and `p` columns.

        @return: structured numpy array
        """

        if not datatype or datatype not in DATATYPES:
            datatype = "TPI"

        tpi = create_data(self._expansion.shape, datatype=datatype, dtype=dtype)
        tpi['t'] = self._scan['t']
        tpi['p'] = self._scan['p']
        tpi['i'] = self._expansion

        return tpi

    @property
    def lmax(self):
        """
        Maximum l angular number used

        Determines the size of the alm array and number of coefficients calculated by `generate`.

        The useful range for XPD scans is from 30 to 90.
        The default is 60, which is sufficient to reproduce fine backscattering details.
        """

        return self._lmax

    @lmax.setter
    def lmax(self, value):
        if value != self._lmax:
            self._lmax = value
            self._alm = np.zeros((self._lmax // 2 + 1, 2 * self._lmax + 1))
            self._sph_harm = None
            self._d_omega = None

    @property
    def alm(self):
        """
        Multipole coefficients a_lm

        Multipole coefficients a_lm are stored in a two-dimensional ndarray.
        The shape is (lmax / 2 + 1, 2 * lmax + 1),
        where the first dimension is l (0...lmax, step 2),
        and the second dimension is m (-lmax...+lmax, step 1).

        There are two ways to modify this array:
        by assignment of a whole array or modification of individual values.

        On assignment of a whole array, lmax is calculated from the shape of the array.
        The array is copied, NaNs are converted to zeros.
        Items |m| > l are ignored and should be zero.
        """

        return self._alm

    @alm.setter
    def alm(self, value):
        lmax = (value.shape[0] - 1) * 2
        assert value.shape == (lmax // 2 + 1, 2 * lmax + 1)
        if lmax != self._lmax:
            self._lmax = lmax
            self._sph_harm = None
            self._d_omega = None
        self._alm = np.nan_to_num(value, copy=True, nan=0.0)

    def generate(self):
        """
        Calculate the multipole coefficients for the stored holoscan

        Requires: lmax and holoscan
        Generates: alm

        @return alm: Multipole coefficients in a two-dimensional ndarray.
            The shape is (lmax / 2 + 1, 2 * lmax + 1),
            where the first dimension is l (0...lmax, step 2),
            and the second dimension is m (-lmax...+lmax, step 1).
            Items |m| > l are 0.
        """

        if not self._check_basis():
            self._calc_basis()

        self._alm = np.sum(self._original * np.conj(self._sph_harm) * self._d_omega, axis=-1)

        return self._alm

    def expand(self, lmin: Optional[int] = None, lmax: Optional[int] = None) -> np.ndarray:
        """
        Calculate the expanded pattern from coefficients

        The alm property and the `t` and `p` columns of `holoscan` must be valid.
        The result is available from `expansion` and as the function result.

        The l range used in the expansion can be constricted by the lmin and lmax options.
        By default, the full alm matrix is used.

        @param lmin: minimum l to include in expansion.
            Allowed range: 0...self.lmax.
        @param lmax: maximum l to include in expansion.
            Allowed range: 0...self.lmax.
            By default, all coefficients are included.
        @return: intensity array versus `holoscan`.
            This array contains just the intensity.
            The `expansion` property provides the full scan object in TPI structure.
            Other structures can be retrieved from the `get_expansion` method.
        """

        if not self._check_basis():
            self._calc_basis()

        if lmax or lmin:
            alm = self._alm.copy()
            if lmin is not None and 0 <= lmin <= self._lmax:
                alm[0:lmin, :] = 0
            if lmax is not None and 0 <= lmax <= self._lmax:
                alm[lmax+1:self._lmax+1, :] = 0
        else:
            alm = self._alm

        self._expansion = np.real(np.sum(np.sum(alm * self._sph_harm.transpose((2, 0, 1)), axis=-1), axis=-1))

        return self._expansion

    def export_alm_h5(self, filename: os.PathLike):
        """
        Export the alm coefficients to an HDF5 file.

        The HDF5 file contains 3 datasets: alm, l and m.
        The alm matrix is written directly as a 2D complex array
        (compound type with members r and i containing the real and imaginary parts, respectively).
        l and m ranges are linked as dimension scales.

        @param filename: path-like
        @return: None
        """

        range_l = np.arange(0, self._lmax + 1, 2)
        range_m = np.arange(-self._lmax, self._lmax + 1, 1)
        with h5py.File(filename, "w") as h5:
            h5["l"] = range_l
            h5["m"] = range_m
            h5["l"].make_scale()
            h5["m"].make_scale()
            h5["alm"] = self._alm
            h5["alm"].dims[0].label = "l"
            h5["alm"].dims[1].label = "m"
            h5["alm"].dims[0].attach_scale(h5["l"])
            h5["alm"].dims[1].attach_scale(h5["m"])

    def import_alm_h5(self, filename: os.PathLike):
        """
        Import the alm coefficients from an HDF5 file.

        The file must use the same structure as created by `export_alm_h5`.
        Alternatively, the real and imaginary parts can be in separate datasets "alm_r" and "alm_i", respectively.

        @param filename: path-like
        @return:
        """

        with h5py.File(filename, "r") as h5:
            try:
                self.alm = h5["alm"][()]
            except KeyError:
                self.alm = complex(h5["alm_r"][()], h5["alm_i"][()])

    def _check_basis(self):
        """
        Check whether the internal basis functions are consistent with scan data.

        # sph dimensions: l, m, tp

        @return:
        """

        nl = self._lmax // 2 + 1
        nm = 2 * self._lmax + 1
        if self._sph_harm is None or self._d_omega is None:
            return False
        if self._sph_harm.shape[0] != nl or self._d_omega.shape[0] != nl:
            return False
        if self._sph_harm.shape[1] != nm or self._d_omega.shape[1] != nm:
            return False
        if self._sph_harm.shape[2] != self._scan.shape[0] or self._d_omega.shape[2] != self._scan.shape[0]:
            return False
        return True

    def _discard_basis(self):
        """
        Discard the basis functions (spherical harmonics and volume element).

        This is necessary when the scan positions, lmax or shape of alm coefficients have changed.
        It is not necessary after changes to the intensity column or alm values.
        The basis functions will be recalculated automatically.

        @return: None
        """

        self._sph_harm = None
        self._d_omega = None

    def _calc_basis(self):
        """
        Calculate the basis functions.

        The basis functions include the spherical harmonics and the volume element.
        They depend on the scan grid and lmax.
        For better performance, the values of the basis functions are evaluated on a particular scan grid
        and stored for further use.

        They have to be re-calculated whenever lmax or the scan angles have changed.
        It is not necessary after changes to the intensity column or alm values.

        @return: None
        """

        range_l = np.arange(0, self._lmax + 1, 2)
        range_m = np.arange(-self._lmax, self._lmax + 1, 1)
        tpi = np.arange(self._scan.shape[0])
        # order of dimensions is l, m, tp
        mesh_m, mesh_l, mesh_tpi = np.meshgrid(range_m, range_l, tpi)
        mesh_theta = np.deg2rad(self._scan['t'][mesh_tpi])
        mesh_phi = np.deg2rad(self._scan['p'][mesh_tpi])

        # sph_harm returns nan for bad l/m combinations
        self._sph_harm = scipy.special.sph_harm(mesh_m, mesh_l, mesh_phi, mesh_theta)
        self._sph_harm = np.nan_to_num(self._sph_harm, copy=False, nan=0.0)

        unique_theta, d_theta, d_phi = analyse_holoscan_steps(self._scan)
        mesh_sin = np.interp(mesh_theta, unique_theta, np.sin(np.deg2rad(unique_theta)))
        mesh_dtheta = np.interp(mesh_theta, unique_theta, np.deg2rad(d_theta))
        mesh_dphi = np.interp(mesh_theta, unique_theta, np.deg2rad(d_phi))
        self._d_omega = mesh_sin * mesh_dtheta * mesh_dphi
