""" Classes for handling the waveform modes or data defined on spheres.


Classes
-------
spherical_array: A 2D data-type.
                  Stores and manages two-dimensional data
                  on surfaces of spherical topology.
modes_array: A data-type.
              Handle and work with mode coefficients.
"""

import sys

import h5py
import numpy as np
import math

from waveformtools import dataIO
from waveformtools.dataIO import (
    construct_mode_list,
    get_iteration_numbers_from_keys,
    sort_keys,
)
from spectral.spherical.grids import UniformGrid
from spectral.spherical.swsh import Yslm_vec
from waveformtools.waveformtools import interp_resam_wfs, message
from pathlib import Path
from waveformtools.integrate import fixed_frequency_integrator
from waveformtools.waveformtools import (
    get_starting_angular_frequency as sang_f,
)


""" Units """

G = 6.67428 * 10 ** (-11.0)  # m^3 /kg /s^2
# c = 3.0 * 10**8  # m/s
c = 299_792_458.0  # m/s
# Msun = 1.32712440041e20
Msun = 1.988500 * 10**30.0  # g, SI
muc = c**2 / (G * Msun)  # g, NR to SI
tuc = G * Msun / (c**3)  # s, NR to SI
# dMpc = 3.0857 * 10 ** (22)  # m
dMpc = 3.0856775814913672789139379577965e22


class spherical_array:
    """A class for handling waveforms on a sphere.


    Attributes
    ----------
    label: str
            The label of the waveform data.
    time_axis: 1d array
                The time axis of the data.
    frequency_axis: 1d array
                     The frequency axis if the data
                     is represented in frequency domain.
    grid_info: UniformGrid
                An instance of the `UniformGrid` class.
    data_len: int
               The length of the data along the time axis.

    Methods
    -------
    delta_t:
              Fetch the time stepping `delta_t`.
    to_modes_array:
                     Find the waveform expressed in the
                     SWSH basis.
    boost:
           Boost the waveform.
    supertranslate:
                    Supertranslate the waveform.


    Notes
    -----
    The modes array has the axis in the form
    (flattened_modes_1d, *shape of extra modes, )
    """

    def __init__(
        self,
        label=None,
        time_axis=None,
        frequency_axis=None,
        data=None,
        data_len=None,
        data_dir=None,
        file_name=None,
        grid_info=None,
        spin_weight=-2,
        ell_max=8,
    ):
        self._label = label
        self._data_len = data_len
        # self.base_dir = base_dir  # The base directory containing the
        self._data = data
        self._file_name = file_name
        self._data_dir = data_dir
        self._time_axis = time_axis
        self._frequency_axis = frequency_axis
        self._grid_info = grid_info
        self._spin_weight = spin_weight
        self._ell_max = ell_max

    @property
    def label(self):
        return self._label

    @property
    def ell_max(self):
        return self._ell_max

    @property
    def time_axis(self):
        return self._time_axis

    @property
    def frequency_axis(self):
        return self._frequency_axis

    @property
    def data(self):
        return self._data

    @property
    def data_dir(self):
        return self._data_dir

    @property
    def file_name(self):
        return self._file_name

    @property
    def spin_weight(self):
        return self._spin_weight

    @property
    def grid_info(self):
        return self._grid_info

    def delta_t(self, value=None):
        """Sets and returns the value of time stepping :math:`dt`.

        Parameters
        ----------
        value: float, optional
                The value of :math:`dt`
                to set to the attribute.

        Returns
        -------
        delta_t: float
                  Sets the attribute.
        """

        if not self._delta_t:
            if not value:
                try:
                    self._delta_t = self.time_axis[1] - self.time_axis[0]
                except Exception as ex:
                    message(
                        "Please input the value of `delta_t`"
                        "or supply the `time_axis` to the waveform.",
                        ex,
                    )
            else:
                self._delta_t = value

        return self._delta_t

    @property
    def data_len(self):
        """Returns the length of the time/frequency axis.

        Returns
        -------
        data_len: int
                   The length of the time/frequency axis.
        """

        if self._data_len is None:
            try:
                data_len = len(self.time_axis)

            except Exception as ex:
                data_len = len(self.frequency_axis)
                message(ex)

            self._data_len = data_len

        return self._data_len

    def to_modes_array(self, grid_info=None, spin_weight=None, ell_max=8):
        """Decompose a given spherical array function on a sphere
        into Spin Weighted Spherical Harmonic modes.

        Parameters
        ----------
        spin_weight: int, optional
                      The spin weight of the waveform.
                      It defaults to -2 for a
                      gravitational waveform.
        ell_max: int, optional
                  The maximum value of the :math:`\\ell`
                  polar quantum number. Defaults to 8.
        grid_info: class instance
                    The class instance that contains
                    the properties of the spherical grid.

        Returns
        -------
        waveforms_modes: modes_array
                          An instance of the `modes_array` class
                          containing the decomposed modes.

        Notes
        -----
        1. Assumes that the sphere on which
           this decomposition is carried out is so far out
           that the coordinate system is spherical polar
           on a round sphere.
        2. Assumes that the poper area is the same
           as its co-ordinate area.
        3. Ensure that the label of the
           input spherical array indicates whether
           it is a time domain data or frequency domain data.
        """

        if self.grid_info is None:
            if grid_info is None:
                message(
                    "Please specify the grid specs. Assuming a default"
                    " uniform grid.",
                    message_verbosity=1,
                )
                grid_info = UniformGrid()

            self._grid_info = grid_info

        if self._ell_max is None:
            if ell_max is None:
                try:
                    self._ell_max = grid_info.L + 1
                except Exception as ex:
                    message(
                        ex,
                        "\n ell_max not provided. Assuming 8",
                        message_verbosity=1,
                    )
                    self._ell_max = 8

            else:
                try:
                    ell_max_grid = grid_info.L + 1

                    assert ell_max <= ell_max_grid, (
                        "This GL grid"
                        "can only decompose upto"
                        f"ell_max of {ell_max_grid}"
                    )

                    self._ell_max = ell_max

                except Exception as ex:
                    message(
                        ex,
                        "No ell_max constraints" "due to uniform grid",
                        message_verbosity=3,
                    )

        if spin_weight is None:
            if self.spin_weight is None:
                message("Please specify spin weight. Assuming 2")
                spin_weight = 2
                self.spin_weight = spin_weight

            else:
                spin_weight = self.spin_weight

        spin_weight = abs(spin_weight)
        # Compute the meshgrid for theta and phi.
        theta, phi = self.grid_info.meshgrid

        # Create a modes array object

        # Create a modes list.
        modes_list = construct_mode_list(ell_max, spin_weight=spin_weight)

        if not self.label:
            label = "decomposed_time_domain"
        else:
            label = self.label

        # Create a mode array for the decomposed_waveform
        waveform_modes = modes_array(
            label=label, ell_max=ell_max, spin_weight=spin_weight
        )

        # Inherit the time or frequency axis.
        if "time" in label:
            # axis = self.time_axis
            waveform_modes.time_axis = self.time_axis
        else:
            # axis = self.frequency_axis
            waveform_modes.frequency_axis = self.frequency_axis

        # Create the modes_array
        waveform_modes.create_modes_array(
            ell_max=ell_max, data_len=self.data_len
        )
        waveform_modes.modes_list = modes_list
        # The area element on the sphere
        # Compute the meshgrid for theta and phi.
        theta, phi = self.grid_info.meshgrid

        # sqrt_met_det = np.sin(theta)
        # sqrt_met_det = np.sqrt(np.power(np.sin(theta), 2))

        # darea = sqrt_met_det * grid_info.dtheta * grid_info.dphi

        modes_list = [item for item in modes_list if item[0] >= spin_weight]

        from waveformtools.integrate import TwoDIntegral

        for mode in modes_list:
            ell_value, all_emm_values = mode

            for emm_value in all_emm_values:
                # m value.

                # Spin weighted spherical harmonic function at (theta, phi)

                Ybasis_fun = np.conj(
                    Yslm_vec(
                        spin_weight,
                        ell=ell_value,
                        emm=emm_value,
                        theta_grid=theta,
                        phi_grid=phi,
                    )
                )

                # Ydarea = Ybasis_fun * darea
                # print(Ydarea.shape)
                # print(full_integrand)
                # Using quad

                # multipole_ell_emm = np.tensordot(
                #    self.data, Ydarea, axes=((0, 1), (0, 1))
                # )

                # print("Shape", Ybasis_fun.shape, self.data.shape)
                # integrand = np.tensordot(self.data,
                # Ybasis_fun, axes=((0, 1), (0, 1)))
                integrand = np.transpose(
                    np.transpose(self.data, (2, 0, 1)) * Ybasis_fun, (1, 2, 0)
                )

                # integrand = self.data * Ybasis_fun
                multipole_ell_emm = TwoDIntegral(integrand, self.grid_info)

                # print(f'l{ell_value}m{emm_value}', multipole_ell_emm)

                waveform_modes.set_mode_data(
                    ell_value, emm_value, multipole_ell_emm
                )

        return waveform_modes

    def boost(self, conformal_factor, grid_info=None):
        """Boost the waveform given the unboosted
        waveform and the boost conformal factor.

        Parameters
        ----------
        self: spherical_array
               A class instance of `spherical array`.

        conformal_factor: 2d array
                           The conformal factor for the
                           Lorentz transformation.
                           It may be a single floating point number
                           or an array on a spherical grid. The array
                           will be of dimensions [ntheta, nphi]

        grid_info: class instance
                    The class instance that contains
                    the properties of the spherical grid.


        Returns
        -------
        boosted_waveform: sp_array
                           The class instance `sp_array` that
                           contains the boosted waveform.
        """

        from waveformtools.waveforms import spherical_array

        if grid_info is None:
            grid_info = self.grid_info

        # Compute the boosted waveform
        # on the spherical grid
        # on all the elements.

        boosted_waveform_data = (
            np.transpose(self.data, (2, 0, 1)) * conformal_factor
        )
        # boosted_waveform_data = np.multiply(self.data,
        # conformal_factor[:, :, None])
        # boosted_waveform_data = boosted_waveform_item
        # boosted_waveform_data = np.array([np.transpose(item)
        # *conformal_factor
        # for item in np.transpose(self.data)])

        # Construct a 2d waveform array object
        boosted_waveform = spherical_array(
            grid_info=grid_info,
            data=np.transpose(np.array(boosted_waveform_data), (1, 2, 0)),
        )

        # Assign other attributes.
        boosted_waveform.label = "boosted " + self.label
        boosted_waveform.time_axis = self.time_axis

        return boosted_waveform

    def supertranslate(self, supertransl_alpha_sp, order=1):
        """Supertranslate the :math:`\\Psi_{4\\ell m}` waveform modes,
        given the, the supertranslation parameter and the order.

        Parameters
        ----------
        supertransl_alpha_modes: modes_array
                                  The modes_array containing the
                                  supertranslation mode coefficients.
        grid_info: class instance
                    The class instance that contains
                    the properties of the spherical grid
                    using which the computations are
                    carried out.
        order: int
                The number of terms to use to
                approximate the supertranslation.

        Returns
        -------
        Psi4_supertransl: modes_array
                           A class instance containg the
                           modes of the supertranslated
                           waveform:math:`\\Psi_4`.
        """

        # Create a spherical_array to hold the supertranslated waveform
        Psi4_supertransl_sp = spherical_array(
            grid_info=self.grid_info,
            label="{} -> supertranslated time".format(self.label),
        )

        delta_t = float(self.delta_t())

        # Set the data.
        data = 0
        # data = self.data
        # Psi4_supertransl_sp.data = self.data
        # delta = 0
        # count = 0
        from waveformtools.differentiate import differentiate5_vec_numba

        for index in range(order):
            # print(f'{index} loop')
            dPsidu = self.data
            for dorder in range(index + 1):
                # print(f'differentiating {dorder+1} times')
                dPsidu = differentiate5_vec_numba(dPsidu, delta_t)

            message("Incrementing...")
            # delta = np.power(supertransl_alpha_sp.data, index+1)
            # * dPsidu / np.math.factorial(index+1)
            # print(delta/self.data)

            data += (
                np.power(supertransl_alpha_sp.data, index + 1)
                * dPsidu
                / math.factorial(index + 1)
            )  # delta

        data += self.data
        message("Equal to original waveform?", (data == self.data).all())

        Psi4_supertransl_sp.data = data
        Psi4_supertransl_sp.time_axis = self.time_axis
        message("Done.")
        return Psi4_supertransl_sp

    def load_qlm_data(
        self, data_dir=None, grid_info=None, bh=0, variable="sigma"
    ):
        """Load the 2D shear data from h5 files.

        Parameters
        ----------
        file_name: str
                    The name of the file containing data.
        data_dir: str
                   The name of the directory containing data.
        grid_info: class instance
                    An instance of the grid_info class.
        bh: int
             The black hole number (0, 1 or 2)
        """

        if data_dir is None:
            if self.data_dir is None:
                print("Please supply the data directory!")
            else:
                data_dir = self.data_dir
        else:
            if self.data_dir is None:
                self.data_dir = data_dir

        if grid_info is None:
            if self.grid_info is None:
                message("Please supply the grid spec!")
            else:
                grid_info = self.grid_info
        else:
            if self.grid_info is None:
                self.grid_info = grid_info
        # get the full path.

        file_name = f"qlm_{variable}[{bh}].xy.h5"

        full_path = self.data_dir + file_name

        # cflag = 0

        nghosts = grid_info.nghosts
        ntheta = grid_info.ntheta
        nphi = grid_info.nphi

        # Open the modes file.
        with h5py.File(full_path, "r") as wfile:
            # Get all the mode keys.
            modes_keys_list = list(wfile.keys())
            # modes_keys_list = sorted(modes_keys_list)

            # Get the mode keys containing the data.
            modes_keys_list = [
                item for item in modes_keys_list if "it=" in item
            ]

            # Get the itaration numbers.
            iteration_numbers = sorted(
                get_iteration_numbers_from_keys(modes_keys_list)
            )
            # sargs = np.argsort(iteration_numbers)
            # iteration_numbers = iteration_numbers[sargs]
            modes_keys_list = sort_keys(modes_keys_list)
            # Construct the data array.

            data_array = []

            for key in modes_keys_list:
                # data_item = np.array(wfile[key])
                # print(data_item.shape)
                data_item = np.array(wfile[key])[
                    nghosts : nphi - nghosts, nghosts : ntheta - nghosts
                ]
                try:
                    data_item = data_item["real"] + 1j * data_item["imag"]
                except Exception as ex:
                    message(ex)
                    pass

                data_array.append(data_item)

        self.data = np.transpose(np.array(data_array), (2, 1, 0))

        self.iteration_axis = np.array(iteration_numbers)

        #########################################################
        # Load inv_coords data
        #########################################################

        inv_file_name = f"qlm_inv_z[{bh}].xy.h5"

        # get the full path.
        full_path = self.data_dir + inv_file_name

        # Open the modes file.
        with h5py.File(full_path, "r") as wfile:
            # Get all the mode keys.
            modes_keys_list = list(wfile.keys())
            # modes_keys_list = sorted(modes_keys_list)

            # Get the mode keys containing the data.
            modes_keys_list = [
                item for item in modes_keys_list if "it=" in item
            ]

            modes_keys_list = sort_keys(modes_keys_list)
            data_array = []

            for key in modes_keys_list:
                data_item = np.array(wfile[key])[
                    nghosts : nphi - nghosts, nghosts : ntheta - nghosts
                ]
                # data_item = data_item['real'] + 1j*data_item['imag']
                data_array.append(data_item)

        self.invariant_coordinates_data = np.transpose(
            np.array(data_array), (2, 1, 0)
        )

        #########################################################
        # Load metric determinant  data
        #########################################################

        twometric_qtt_file_name = f"qlm_qtt[{bh}].xy.h5"
        twometric_qtp_file_name = f"qlm_qtp[{bh}].xy.h5"
        twometric_qpp_file_name = f"qlm_qpp[{bh}].xy.h5"

        # set the full path.
        full_path = self.data_dir + twometric_qtt_file_name

        # Open the modes file.
        with h5py.File(full_path, "r") as wfile:
            # Get all the mode keys.
            modes_keys_list = list(wfile.keys())
            # modes_keys_list = sorted(modes_keys_list)

            # Get the mode keys containing the data.
            modes_keys_list = [
                item for item in modes_keys_list if "it=" in item
            ]

            modes_keys_list = sort_keys(modes_keys_list)

            qtt_data_array = []

            for key in modes_keys_list:
                data_item = np.array(wfile[key])[
                    nghosts : nphi - nghosts, nghosts : ntheta - nghosts
                ]
                # data_item = data_item['real'] + 1j*data_item['imag']
                qtt_data_array.append(data_item)

        qtt_data_array = np.array(qtt_data_array)
        qtt_data_array = np.transpose(qtt_data_array, (2, 1, 0))

        # set the full path.
        full_path = self.data_dir + twometric_qtp_file_name

        # Open the modes file.
        with h5py.File(full_path, "r") as wfile:
            # Get all the mode keys.
            modes_keys_list = list(wfile.keys())
            # modes_keys_list = sorted(modes_keys_list)

            # Get the mode keys containing the data.
            modes_keys_list = [
                item for item in modes_keys_list if "it=" in item
            ]

            modes_keys_list = sort_keys(modes_keys_list)

            qtp_data_array = []

            for key in modes_keys_list:
                data_item = np.array(wfile[key])[
                    nghosts : nphi - nghosts, nghosts : ntheta - nghosts
                ]
                # data_item = data_item['real'] + 1j*data_item['imag']
                qtp_data_array.append(data_item)

        qtp_data_array = np.array(qtp_data_array)
        qtp_data_array = np.transpose(qtp_data_array, (2, 1, 0))

        # set the full path.
        full_path = self.data_dir + twometric_qpp_file_name

        # Open the modes file.
        with h5py.File(full_path, "r") as wfile:
            # Get all the mode keys.
            modes_keys_list = list(wfile.keys())
            # modes_keys_list = sorted(modes_keys_list)

            # Get the mode keys containing the data.
            modes_keys_list = [
                item for item in modes_keys_list if "it=" in item
            ]

            modes_keys_list = sort_keys(modes_keys_list)

            qpp_data_array = []

            for key in modes_keys_list:
                data_item = np.array(wfile[key])[
                    nghosts : nphi - nghosts, nghosts : ntheta - nghosts
                ]
                # data_item = data_item['real'] + 1j*data_item['imag']
                qpp_data_array.append(data_item)

        qpp_data_array = np.array(qpp_data_array)
        qpp_data_array = np.transpose(qpp_data_array, (2, 1, 0))

        sqrt_met_det = np.sqrt(
            np.linalg.det(
                np.transpose(
                    np.array(
                        [
                            [qtt_data_array, qtp_data_array],
                            [qtp_data_array, qpp_data_array],
                        ]
                    ),
                    (2, 3, 4, 0, 1),
                )
            )
        )

        self.sqrt_met_det_data = sqrt_met_det

    def to_shear_modes_array(
        self, grid_info=None, spin_weight=None, ell_max=8
    ):
        """Decompose a given spherical array function on a sphere
        into Spin Weighted Spherical Harmonic modes.

        Parameters
        ----------
        spin_weight: int, optional
                      The spin weight of the waveform.
                      It defaults to -2 for
                      a gravitational waveform.
        ell_max: int, optional
                  The maximum value of the :math:`\\ell`
                  polar quantum number. Defaults to 8.
        grid_info: class instance
                    The class instance that contains
                    the properties of the spherical grid.

        Returns
        -------
        waveforms_modes: modes_array
                          An instance of the `modes_array`
                          class containing the decomposed modes.

        Notes
        -----
        1. Assumes that the sphere on which
           this decomposition is carried out is so far out
           that the coordinate system is spherical polar
           on a round sphere.
        2. Assumes that the poper area is
           the same as its co-ordinate area.
        3. Ensure that the label of the input
           spherical array indicates whether
           it is a time domain data or
           frequency domain data.
        """

        if grid_info is None:
            if self.grid_info is None:
                message("Please specify the grid specs. Assuming defaults.")
                grid_info = UniformGrid()
                self.grid_info = grid_info
            else:
                grid_info = self.grid_info

        if spin_weight is None:
            if self.spin_weight is None:
                message("Please specify spin weight. Assuming 2")
                spin_weight = 2
                self.spin_weight = spin_weight

            else:
                spin_weight = self.spin_weight

        spin_weight = abs(spin_weight)
        # Compute the meshgrid for theta and phi.
        theta, phi = grid_info.meshgrid

        # Create a modes array object

        # Create a modes list.
        modes_list = construct_mode_list(ell_max, spin_weight=spin_weight)

        if not self.label:
            label = "decomposed_time_domain"
        else:
            label = self.label

        # Create a mode array for the decomposed_waveform
        waveform_modes = modes_array(
            label=label, ell_max=ell_max, spin_weight=spin_weight
        )

        # Inherit the time or frequency axis.
        if "time" in label:
            # axis = self.time_axis
            waveform_modes.time_axis = self.time_axis
        else:
            # axis = self.frequency_axis
            waveform_modes.frequency_axis = self.frequency_axis

        # Create the modes_array
        waveform_modes.time_axis = self.time_axis[:]
        # sargs = np.argsort(waveform_modes.time_axis)
        # message(sargs)
        waveform_modes.time_axis = waveform_modes.time_axis

        waveform_modes.create_modes_array(
            ell_max=ell_max, data_len=self.data_len
        )
        waveform_modes.modes_list = modes_list
        # The area element on the sphere
        # Compute the meshgrid for theta and phi.
        theta, phi = grid_info.meshgrid

        phi = np.transpose(
            np.array([phi for index in range(len(self.time_axis))]), (1, 2, 0)
        )

        # sqrt_met_det = np.sin(theta)
        # sqrt_met_det = np.sqrt(np.power(np.sin(theta), 2))

        darea = self.sqrt_met_det_data * grid_info.dtheta * grid_info.dphi

        theta = np.emath.arccos(self.invariant_coordinates_data)

        modes_list = [item for item in modes_list if item[0] >= spin_weight]

        for mode in modes_list:
            ell_value, all_emm_values = mode

            for emm_value in all_emm_values:
                # m value.
                # message(f'Processing l{ell_value} m{emm_value}')
                # Spin weighted spherical harmonic function at (theta, phi)

                Ybasis_fun = np.conj(
                    Yslm_vec(
                        spin_weight=spin_weight,
                        ell=ell_value,
                        emm=emm_value,
                        theta_grid=theta,
                        phi_grid=phi,
                    )
                )
                # Ybasis_fun = np.array([np.conj(Yslm_vec(spin_weight=
                # spin_weight, ell=ell_value, emm=emm_value,
                # theta_grid=theta[:, :, index],
                # phi_grid=phi[:, :, index])) for index in
                # range(self.data_len)])
                # Ybasis_fun = np.transpose(Ybasis_fun, (1, 2, 0))
                # message('Ybasis_fun', Ybasis_fun.shape)
                Ydarea = Ybasis_fun * darea
                # message('Ydarea', Ydarea.shape)
                # message(full_integrand)
                # Using quad
                # message('self.data', self.data.shape)
                # multipole_ell_emm = np.tensordot(self.data, Ydarea,
                # axes=((0, 1), (0, 1)))
                multipole_ell_emm = np.sum(self.data * Ydarea, (0, 1))

                # message(f'l{ell_value}m{emm_value}', multipole_ell_emm)

                # message('multipole_ell_emm', multipole_ell_emm.shape)
                waveform_modes.set_mode_data(
                    ell_value, emm_value, data=multipole_ell_emm
                )

        return waveform_modes


# @jitclass(spec_ma)
class modes_array:
    """A class that holds mode array of waveforms

    This can handle two index and three index modes.

    Attributes
    ----------
    label: str
            The label of the modes array.
    r_ext: float
            The extraction radius.
    modes_list: list
                 The list of available modes
                 in the format [l1, [m values], [l2, [m values], ...]]
    ell_max: int
              The maximum :math:`\\ell` mode available.
    modes_data: 3d array
                 The three dimensional array
                 containing modes in time/frequency
                 space. The axis of the array is
                 (:math:`\\ell`, :math:`m`, time/freq).
    base_dir: str
               The base directory containing the
               modes data.
    data_dir: str
               The subdirectory in which to look
               for the data.
    filename: str
               The filename containg the modes data.

    Methods
    -------
    get_metadata:
                   Get the metadata associated with the modes_array.
    mode:
           Get the data for the given :math:`\\ell, m` mode.
    create_modes_array:
                         A private method to create an
                         empty `modes_array` of given shape.
    delta_t:
              Set the attribute `delta_t` and/ or return its value.
    load_modes:
                Load the waveform modes from a specified h5 file.
    save_modes:
                Save the waveform modes to a specified h5 file.
    set_mode_data:
                   Set the `mode` data of specified modes.
    to_frequency_basis:
                        Get the `modes_array` in frequency basis
                        from its time basis representation.
    to_time_basis:
                   Get the `modes_array` in temporal basis
                   from its frequency basis representation.
    extrap_to_inf:
                   Extrapolate the modes to infinity.
    supertranslate:
                    Supertranslate the waveform modes.
    boost:
           Boost the waveform modes.
    """

    def __init__(
        self,
        data_dir=None,
        file_name=None,
        extra_mode_axes_shape=None,
        modes_data=None,
        time_axis=None,
        frequency_axis=None,
        key_format=None,
        ell_max=None,
        data_len=None,
        modes_list=None,
        label=None,
        r_ext=np.inf,
        out_file_name=None,
        maxtime=None,
        date=None,
        time=None,
        key_ex=None,
        spin_weight=-2,
        actions="empty",
        areal_radii=[],
    ):
        self.label = label
        self.data_dir = data_dir
        self.file_name = file_name
        self.modes_data = modes_data
        self.key_format = key_format
        self.ell_max = ell_max
        self.modes_list = modes_list
        self.r_ext = r_ext
        self.time_axis = time_axis
        self.frequency_axis = frequency_axis
        self.out_file_name = out_file_name
        self.maxtime = maxtime
        self.date = date
        self.time = time
        self.key_ex = key_ex
        self._data_len = data_len
        self._spin_weight = spin_weight
        self._actions = actions
        self._extra_mode_axes_shape = extra_mode_axes_shape
        self._areal_radii = areal_radii

        if (np.array(self.extra_mode_axes_shape) == np.array(None)).all():
            self.extra_mode_axes = False
        else:
            self.extra_mode_axes = True

    @property
    def extra_mode_axes_shape(self):
        return self._extra_mode_axes_shape

    @property
    def actions(self):
        return self._actions

    @property
    def spin_weight(self):
        return self._spin_weight

    def get_metadata(self):
        """Get the metadata associated with the instance.

        Returns
        -------
        metadata: dict
                   A dictionary of metedata.
        """
        # The metadata dict
        unnecessary_keys = ["_time_axis", "_modes_data", "freq_axis"]

        # Get all attributes
        # metadata = self.__dict__
        metadata = {}

        for key, val in self.__dict__.items():
            if key in unnecessary_keys:
                pass
            else:
                metadata.update({key: val})

        # for item in unnecessary_keys:
        # metadata.pop(item, None)

        # self.metadata = metadata

        return metadata

    def mode(self, ell, emm, extra_indices=None):
        """Return the time series of a particular mode.

        Parameters
        ----------
        ell: int
              The :math:`\\ell` index of the mode.
        emm: int
              The :math:`m` index of the mode.
        r_index: int, optional
                  The third mode axis index, if any

        Returns
        -------
        mode_data: array
                    The array of the requested mode.
        """
        vec_index = ell**2 + emm + ell

        ang_mode_data = self.modes_data[vec_index]

        if self.extra_mode_axes:
            if not (np.array(extra_indices) == np.array(None)).all():
                message(f"Extra indices supplied {extra_indices}...")
                return ang_mode_data[*extra_indices]

        return ang_mode_data

    @property
    def time_axis(self):
        """The time axis"""
        return np.array(self._time_axis)

    @time_axis.setter
    def time_axis(self, time_axis):
        self._time_axis = time_axis

    @property
    def modes_data(self):
        """The modes array"""
        return np.array(self._modes_data)

    @modes_data.setter
    def modes_data(self, modes_data):
        self._modes_data = modes_data

    def create_time_axis(self, data_len):
        message(f"Created time axis of length {data_len}", message_verbosity=3)
        self._time_axis = np.zeros(data_len)

    def create_modes_array(self, ell_max=None, data_len=None):
        """Create a modes array and initialize it with zeros.

        Parameters
        ----------
        ell_max: int
                  The maximum :math:`\\ell` value of the modes.
        data_len: int
                   The number of points along
                   the third (time / frequency) axis.

        Returns
        -------
        self.modes_array: modes_array
                           sets the `self.modes_array` attribute.
        """
        import datetime
        import time

        # Check ell_max
        if ell_max is None:
            try:
                ell_max = self.ell_max
            except Exception as ex:
                message(ex)
                raise NameError("Please supply ell_max")

        if data_len is None:
            try:
                data_len = self.data_len
            except Exception as ex:
                message(ex)
                raise NameError("Please supply data_len")

        if (np.array(self.time_axis) == np.array(None)).all():
            self.create_time_axis(data_len)

        if self.modes_list is None:
            self.modes_list = construct_mode_list(
                ell_max=ell_max, spin_weight=self.spin_weight
            )

        if self.extra_mode_axes:
            message(
                "Creating modes array"
                " with extra mode axis"
                f" of shape {self.extra_mode_axes_shape}",
                message_verbosity=3,
            )

            self._modes_data = np.zeros(
                (
                    (ell_max + 1) ** 2,
                    *self.extra_mode_axes_shape,
                    data_len,
                ),
                dtype=np.complex128,
            )

        else:
            message("Creating modes array", message_verbosity=3)

            self._modes_data = np.zeros(
                (
                    (ell_max + 1) ** 2,
                    data_len,
                ),
                dtype=np.complex128,
            )

        # Set the time metadata
        time_now = time.localtime()
        time_now = time.strftime("%H:%M:%S", time_now)

        date_now = str(datetime.date.today())

        if self.time is None:
            # Assign time and date stamp if it doesnt exist
            self.time = time_now
            self.date = date_now

    @property
    def data_len(self):
        """Returns the length of the time/frequency axis.

        Returns
        -------
        data_len: int
                   The length of the time/frequency axis.
        """

        if self._data_len is None:
            try:
                data_len = len(self.time_axis)

            except Exception as ex:
                message(ex)
                data_len = len(self.frequency_axis)

            self._data_len = data_len

        return self._data_len

    @data_len.setter
    def data_len(self, data_len):
        self._data_len = data_len

    def delta_t(self, value=None):
        """Sets and returns the value of time stepping :math:`dt`.

        Parameters
        ----------
        value: float, optional
                The value of :math:`dt`
                to set to the attribute.

        Returns
        -------
        self.delta_t: float
                       Sets the attribute.
        """

        if not value:
            try:
                delta_t = self.time_axis[1] - self.time_axis[0]
            except Exception as ex:
                message(
                    "Please input the value of `delta_t`"
                    "or supply the `time_axis` to the waveform.",
                    ex,
                )
        else:
            delta_t = value

        return delta_t

    @property
    def delta_f(self, value=None):
        """Sets and returns the value of frequency stepping :math:`df`.

        Parameters
        ----------
        value: float, optional
                The value of :math:`df`
                to set to the attribute.

        Returns
        -------
        delta_f: float
                  Sets the attribute.
        """

        # if not self.delta_t:
        if not value:
            try:
                delta_f = self.frequency_axis[1] - self.frequency_axis[0]
            except Exception as ex:
                message(
                    "Please input the value of `delta_f`"
                    "or supply the `frequency_axis` to the waveform.",
                    ex,
                )
        else:
            delta_f = value

        return delta_f

    def load_modes(
        self,
        label=None,
        data_dir=None,
        file_name=None,
        ftype="generic",
        var_type="Psi4",
        resam_type="finest",
        interp_kind="cubic",
        extrap_order=4,
        r_ext=None,
        ell_max=None,
        pre_key=None,
        modes_list=None,
        crop=False,
        centre=True,
        key_ex=None,
        save_as_ma=False,
        compression_opts=None,
        r_ext_factor=1,
        debug=False,
    ):
        """Load the waveform mode data from an hdf file.

        Parameters
        ----------
        extrap_order: int, optional
                       For SpEC waveforms only.
                       This is the order of extrapolation to use.

        pre_key: str, optional
                  A string containing the key of the group in
                  the HDF file in which the modes` dataset exists.
                  It defaults to `None`.
        mode_numbers: list
                       The mode numbers to load from the file.
                       Each item in the list is a list that
                       contains two integrer numbers, one for
                       the mode index :math:`\\ell` and the
                       other for the mode index :math:`m`.
        crop: bool
               Whether or not to crop the beginning of the input
               waveform. If yes, the first :math:`t=r_{ext}`
               length of data will be discarded.

        Returns
        -------
        waveform_obj: 3d array
                       Sets the three dimensional array
                       `waveform.modes` that contains
                       the required :math:`\\ell, m` modes.

        Examples
        --------
        # Note
        # Update this!
        #>>> from waveformtools.waveforms import waveform
        #>>> waveform.set_basedir('./')
        #>>> waveform.set_datadir('data/')
        #>>> mode_numbers = [[2, 2], [3, 3]]
        #>>> waveform.load_data(mode_numbers=mode_numbers)
        """

        if debug is False:
            wfs_nl = 1

        if not data_dir:
            data_dir = self.data_dir
        else:
            self.data_dir = data_dir

        data_dir = Path(data_dir)

        if not file_name:
            file_name = self.file_name
        else:
            self.file_name = file_name

        if not ell_max:
            ell_max = self.ell_max
        else:
            self.ell_max = ell_max

        if not label:
            label = self.label

        # if self.data_dir is not None:
        # data_dir = self.data_dir

        # if self.file_name is not None:
        # file_name = self.file_name
        message(f"Passing {data_dir}/{file_name}", message_verbosity=3)

        if ftype == "generic":
            dataIO.load_gen_data_from_disk(
                wfa=self,
                label=label,
                data_dir=data_dir,
                file_name=file_name,
                r_ext=r_ext,
                ell_max=ell_max,
                pre_key=pre_key,
                modes_list=modes_list,
                crop=crop,
                centre=centre,
                key_ex=key_ex,
                r_ext_factor=r_ext_factor,
            )

        elif (ftype) == "RIT" or (ftype == "GT"):
            if var_type == "Psi4":
                data_file_path = data_dir / file_name

                dataIO.load_RIT_Psi4_data_from_disk(
                    wfa=self,
                    data_file_path=data_file_path,
                    resam_type=resam_type,
                    interp_kind=interp_kind,
                    ell_max=ell_max,
                    modes_list=modes_list,
                    crop=crop,
                    centre=centre,
                )

            elif var_type == "Strain":
                # message(file_name)
                dataIO.load_RIT_Strain_data_from_disk(
                    self,
                    data_dir=data_dir,
                    file_name=file_name,
                    label=label,
                    resam_type=resam_type,
                    interp_kind=interp_kind,
                    ell_max=ell_max,
                    save_as_ma=save_as_ma,
                    modes_list=modes_list,
                    crop=crop,
                    centre=centre,
                    r_ext_factor=r_ext_factor,
                    debug=debug,
                )
            else:
                message(f"Data {ftype} {var_type} not supported yet!")
                sys.exit(0)

        elif ftype == "SpEC":
            wfs_nl = dataIO.load_SpEC_data_from_disk(
                self,
                label,
                data_dir,
                file_name,
                extrap_order,
                r_ext,
                ell_max,
                centre,
                modes_list,
                save_as_ma,
                resam_type,
                interp_kind,
                compression_opts,
                r_ext_factor,
                debug,
            )

        elif ftype == "SpEC_raw":
            wfs_nl = dataIO.load_SpEC_non_extrap_data_from_disk(
                self,
                label,
                data_dir,
                file_name,
                r_ext,
                ell_max,
                centre,
                modes_list,
                save_as_ma,
                resam_type,
                interp_kind,
                compression_opts,
                r_ext_factor,
                debug,
            )

        elif ftype == "SpECTRE":
            dataIO.load_SpECTRE_data_from_disk(
                self,
                label=label,
                data_dir=data_dir,
                file_name=file_name,
                r_ext=r_ext,
                ell_max=ell_max,
                centre=centre,
                modes_list=modes_list,
                save_as_ma=save_as_ma,
                resam_type=resam_type,
                kind=interp_kind,
                compression_opts=compression_opts,
                r_ext_factor=r_ext_factor,
            )
        else:
            message(f"Data {ftype} {var_type} not supported yet!")
            sys.exit(0)

        return wfs_nl

    def save_modes(
        self,
        ell_max=None,
        pre_key=None,
        key_format=None,
        modes_to_save=None,
        out_file_name="mp_new_modes.h5",
        r_ext_factor=None,
        compression_opts=0,
        r_ext=None,
    ):
        """Save the waveform mode data to an hdf file.

        Parameters
        ----------
        pre_key: str, optional
                  A string containing the key of the group in
                  the HDF file in which the modes` dataset exists.
                  It defaults to `None`.
        mode_numbers: list
                       The mode numbers to load from the file.
                       Each item in the list is a list that
                       contains two integrer numbers, one for
                       the mode index :math:`\\ell` and the
                       other for the mode index :math:`m`.

        Returns
        -------
        waveform_obj: 3d array
                       Sets the three dimensional array `waveform.modes`
                       that contains the required :math:`\\ell, m` modes.

        Examples
        --------
        >>> from waveformtools.waveforms import waveform
        >>> waveform.set_basedir('./')
        >>> waveform.set_datadir('data/')
        >>> mode_numbers = [[2, 2], [3, 3]]
        >>> waveform.load_data(mode_numbers=mode_numbers)
        """
        # import dataIO
        from waveformtools import dataIO

        dataIO.save_modes_data_to_gen(
            self,
            ell_max=None,
            pre_key=None,
            key_format=None,
            modes_to_save=None,
            out_file_name="mp_new_modes.h5",
            r_ext_factor=None,
            compression_opts=0,
            r_ext=None,
        )

    def resample(self, new_time_axis=None, new_delta_t=None):
        """Resample all the waveform modes in time"""

        raise NotImplementedError

    def set_mode_data(
        self, ell_value, emm_value, data, extra_mode_indices=None
    ):
        """Set the mode array data
        for the respective :math:`(\\ell, m)` mode. If the modes array
        has a mode axis of length more then one e.g. if one is dealing
        with not two but three index modes, them one needs to specify
        the third (r) axis index to which this data corresponds to.

        If r_index is not given, then it is assumed that
        the supplied `data` is 2 dimensional (ell, emm, all r)

        Else, only the (ell, emm, one r element) is updated.

        If mode axis length is 1, then it corresponds to usual
        two index modes like (l, m) and they are updated accordingly.

        Parameters
        ----------
        ell_value: int
                    The :math:`\\ell` polar mode number.
        emm_value: int
                    The :math:`emm` azimuthal mode number.
        data: 1d array
               The array consisting of
               mode data for the requested mode.
        extra_mode_indices: tuple
                            A tuple containing the additional
                            mode indices locating the data to
                            be set.

        Returns
        -------
        self.mode_data: modes_data
                         The updated modes data.
        """

        # Compute the emm index given ell.
        emm_index = emm_value + ell_value

        vec_idx = (ell_value) ** 2 + emm_value + ell_value

        if self.extra_mode_axes:
            if extra_mode_indices is not None:
                self._modes_data[vec_idx, *extra_mode_indices] = data
        else:
            # Set the mode data.
            self._modes_data[vec_idx] = data
            # self._modes_data[ell_value, emm_index] = data

    def set_mode_data_at_t_step(
        self, t_step, time_stamp, ell, emm, data, extra_mode_indices=None
    ):
        """Set the mode array data
        for the respective :math:`(\\ell, m)` mode.
        at a single time step.

        Parameters
        ----------
        t_step: int
                 The time step of the mode
        ell_value: int
                    The :math:`\\ell` polar mode number.
        emm_value: int
                    The :math:`emm` azimuthal mode number.
        r_index: int, optional
                  The index of the third mode axis, if any
        data: 1d array
               The array consisting of
               mode data for the requested mode.
        extra_mode_indices: tuple
                            A tuple containing the additional
                            mode indices locating the data to
                            be set.
        Returns
        -------
        self.mode_data: modes_data
                         The updated modes data.
        """
        # Compute the emm index given ell.
        # emm_index = emm + ell
        vec_idx = (ell_value) ** 2 + emm_value + ell_value

        if self.extra_mode_axes:
            if extra_mode_indices is not None:
                message(
                    f"Setting mode data at a mode indices {extra_mode_indices}",
                    message_verbosity=4,
                )

                self._modes_data[vec_idx, *extra_mode_indices, ..., t_step] = (
                    data
                )

        else:
            message(
                "Setting mode data at just ell, emm (no r axis found)",
                message_verbosity=4,
            )
            self._modes_data[vec_idx, ..., t_step] = data

    def to_spherical_array(self, grid_info, meth_info, spin_weight=None):
        """Obtain the spherical array from the modes array.

        Parameters
        ----------
        grid_info: cls instance
                    An instance of the "UniformGrid" class
                    to hold the grid info.
        meth_info: cls instance
                    An instance of the class
                    `waveformtools.diagnostics.method_info` that
                    provides information on
                    what methods to use for integration.

        Returns
        -------
        waveform_sp: spherical_array
                      A member of the "spherical_array" class
                      constructed from the given "modes_rray".
        """

        # Create a spherical array.
        waveform_sp = spherical_array(label=self.label, grid_info=grid_info)

        if spin_weight is None:
            if self.spin_weight is not None:
                spin_weight = self.spin_weight
            else:
                spin_weight = -2
                self.spin_weight = spin_weight

        spin_weight = abs(spin_weight)
        waveform_sp._spin_weight = spin_weight
        # Set the time-axis
        try:
            waveform_sp._time_axis = self.time_axis

        except Exception as ex:
            message(ex)
            waveform_sp._frequency_axis = self.frequency_axis

        # Get the coordinate meshgrid
        theta, phi = grid_info.meshgrid

        s1, s2 = theta.shape
        s3 = self.data_len
        sp_data = np.zeros((s1, s2, s3), dtype=np.complex128)

        modes_list = [
            item for item in self.modes_list if item[0] >= spin_weight
        ]
        for item in modes_list:
            # Get modes.
            ell, emm_list = item
            # if ell<spin_weight:
            # continue

            for emm in emm_list:
                # For every l, m
                sp_data += np.multiply.outer(
                    Yslm_vec(
                        spin_weight,
                        ell=ell,
                        emm=emm,
                        theta_grid=theta,
                        phi_grid=phi,
                    ),
                    self.mode(ell, emm),
                )
                # message(sp_data)
        # Set the data of the spherical array.
        waveform_sp._data = sp_data

        waveform_sp._areal_radii = self._areal_radii

        # try:
        #    waveform_sp._time_axis = self.time_axis
        # except Exception as ex:
        #    message(ex)
        #    waveform_sp._frequency_axis = self.frequency_axis

        return waveform_sp

    def trim(self, trim_upto_time=None):
        """Trim the modes_array at the beginning and center about
        the peak of the 2,2 mode.

        Parameters
        ----------
        time: float
               The time unit upto which to discard.

        Returns
        -------
        Re-sets the `time_axis` and `modes_array` data.
        """
        if trim_upto_time is None:
            trim_upto_time = self.r_ext

        # Compute the start index
        start = int(trim_upto_time / self.delta_t())

        # Trim the time axis
        self._time_axis = self.time_axis[start:]

        # Trim the data
        self._modes_data = self.modes_data[..., start:]

        # Recenter the time axis
        max_ind = np.argmax(np.absolute(self.mode(2, 2)))
        self._time_axis = self.time_axis - self.time_axis[max_ind]

    def to_frequency_basis(self):
        """Compute the modes in frequency basis.

        Returns
        -------
        waveform_tilde_modes: modes_array
                               A modes_array containing the modes
                               in frequency basis.
        """

        # Create a new modes array
        waveform_tilde_modes = modes_array(
            label=f"{self.label} -> frequency_domain"
        )
        waveform_tilde_modes.create_modes_array(
            ell_max=self.ell_max, data_len=self.data_len
        )

        from spectral.fourier.fft import compute_fft

        for mode in self.modes_list:
            # Extrapolate every mode

            # Ge the ell value
            ell_value, emm_list = mode

            for emm_value in emm_list:
                freq_axis, freq_data = compute_fft(
                    self.mode(ell_value, emm_value), self.delta_t()
                )

                waveform_tilde_modes.set_mode_data(
                    ell_value, emm_value, freq_data
                )

        waveform_tilde_modes.frequency_axis = freq_axis
        waveform_tilde_modes.ell_max = self.ell_max
        waveform_tilde_modes.modes_list = self.modes_list
        return waveform_tilde_modes

    def to_time_basis(self):
        """Compute the modes in time basis.

        Returns
        -------
        waveform_modes: modes_array
                         A modes_array containing the modes
                         in frequency basis.
        """

        # Create a new modes array
        waveform_modes = modes_array(label=f"{self.label} -> time_domain")
        waveform_modes.create_modes_array(
            ell_max=self.ell_max, data_len=self.data_len
        )

        from spectral.fourier.fft import compute_ifft

        for mode in self.modes_list:
            # Extrapolate every mode

            # Ge the ell value
            ell_value, emm_list = mode

            for emm_value in emm_list:
                time_axis, time_data = compute_ifft(
                    self.mode(ell_value, emm_value), self.delta_f
                )

                waveform_modes.set_mode_data(ell_value, emm_value, time_data)

        try:
            maxloc = np.argmax(np.absolute(waveform_modes.mode(2, 2)))
        except Exception as ex:
            message(ex)
            maxloc = 0

        maxtime = time_axis[maxloc]

        waveform_modes.time_axis = time_axis - maxtime

        return waveform_modes

    def extrap_to_inf(
        self,
        mass=1,
        spin=None,
        modes_list="all",
        method="SIO",
        r_ext_factor=1,
        diff_method="CS",
        diff_degree=24,
    ):
        """Extrapolate the :math:`\\Psi_4` modes to infinity
        using the perturbative improved second order method.

        Parameters
        ----------
        mass: float
               The effective total mass of the system.
        spin: float
               The effective spin of the system.
        modes: modes array, optional
                The modes to extrapolate. Defaults
                to `all` if not specified.
        method: str
                 The method to use for extrapolation.
                 The available methods are:

        +------------+--------------------------------------+
        | Method str | Name                                 |
        +------------+--------------------------------------+
        |'FO'        | First order                          |
        |'SO'        | Second order                         |
        |'SIO'       | Second improved order                |
        |'NM'        | Numerical method (not ready yet)     |
        +------------+--------------------------------------+

        Returns
        -------
        waveform_inf_modes: modes array
                             A new modes array
                             that contains
                             the extrapolated modes.
        """

        from functools import partial

        # Prepare the extrapolating method.
        if method == "SIO":
            from waveformtools.extrapolate import (
                waveextract_to_inf_perturbative_twop5_order,
            )

            extrap_method = partial(
                waveextract_to_inf_perturbative_twop5_order,
                delta_t=self.delta_t(),
                areal_radius=self.r_ext,
                mass=mass,
                spin=spin,
                method=diff_method,
                degree=diff_degree,
            )

        if method == "SO":
            from waveformtools.extrapolate import (
                waveextract_to_inf_perturbative_two_order,
            )

            extrap_method = partial(
                waveextract_to_inf_perturbative_two_order,
                delta_t=self.delta_t(),
                areal_radius=self.r_ext,
                mass=mass,
                spin=spin,
            )

        if method == "FO":
            from waveformtools.extrapolate import (
                waveextract_to_inf_perturbative_one_order,
            )

            extrap_method = partial(
                waveextract_to_inf_perturbative_one_order,
                u_ret=self.time_axis,
                areal_radius=self.r_ext,
                mass=mass,
            )

        if method == "NM":
            message("This method is not available yet! ")

        # Prepare the modes to be extrapolated.
        if modes_list == "all":
            modes_list = construct_mode_list(self.ell_max, self.spin_weight)

        # Create a mode array for the extrapolated waveform.
        extrap_wf = modes_array(
            label=f"{self.label} -> rPsi4_inf",
            modes_list=self.modes_list,
            ell_max=self.ell_max,
            r_ext=self.r_ext,
        )

        extrap_wf.create_modes_array(
            ell_max=self.ell_max, data_len=self.data_len
        )

        # Retain the time axis.
        extrap_wf.time_axis = self.time_axis
        for mode in modes_list:
            # Extrapolate every mode

            # Ge the ell value
            ell_value, emm_list = mode

            for emm_value in emm_list:
                # For every emm value
                message(f"Processing l{ell_value}, m{emm_value}")
                # Compute rPsi4_lm
                mode_data = r_ext_factor * self.mode(ell_value, emm_value)

                # Extrapolate
                # import ipdb; ipdb.set_trace()
                extrap_mode_data = extrap_method(rPsi4_rlm=mode_data)

                # Assign data to new modes array
                extrap_wf.set_mode_data(ell_value, emm_value, extrap_mode_data)

        message("Done!")
        return extrap_wf

    def supertranslate(self, supertransl_alpha_modes, grid_info, order=4):
        """Supertranslate the :math:`\\Psi_{4\\ell m}`
        waveform modes, give then, the supertranslation parameter
        and the order.

        Parameters
        ----------
        supertransl_alpha_modes: modes_array
                                  The modes_array containing the
                                  supertranslation mode coefficients.
        grid_info: class instance
                    The class instance that contains
                    the properties of the spherical grid
                    using which the computations are carried out.
        order: int
                The number of terms to use
                to approximate the supertranslation.

        Returns
        -------
        Psi4_supertransl: modes_array
                           A class instance containg the
                           modes of the supertranslated
                           waveform :math:`\\Psi_4`.
        """

        import BMS

        ell_max = self.ell_max
        # Step 0: Get the grid properties for integrations

        # Compute the meshgrid for theta and phi.
        theta, phi = grid_info.meshgrid
        # theta
        # Step 1: get the grid function for supertranslation parameter
        supertransl_alpha_sphere = BMS.compute_supertransl_alpha(
            supertransl_alpha_modes, theta, phi
        )

        # The supertranslation is carried out in frequency space.
        # Step 2: get the FFT of the mode coefficients of Psi4
        Psi4_tilde_modes = self.to_frequency_basis()

        # Get the 2d angular frequency array
        omega_axis_2d = Psi4_tilde_modes.omega

        # Construct the supertranslation factor
        supertransl_factor = np.sum(
            np.array(
                [
                    np.power(
                        (-1j * omega_axis_2d * supertransl_alpha_sphere), index
                    )
                    for index in range(order)
                ]
            ),
            axis=0,
        )

        # Multiply with the fourier modes.
        supertransl_spherical_factor = Psi4_tilde_modes.multiply(
            supertransl_factor
        )

        # Reconstruct the modes

        # Check!
        supertransl_spherical_grid = np.zeros(
            supertransl_spherical_factor.shape, dtype=np.complex128
        )

        for ell_value in range(ell_max + 1):
            for emm_value in range(-ell_value, ell_value + 1):
                # Multiply with the SWSH basis.
                supertransl_spherical_grid += (
                    supertransl_spherical_factor
                    * Yslm_vec(
                        spin_weight=-2,
                        ell=ell_value,
                        emm=emm_value,
                        theta=theta,
                        phi=phi,
                    )
                )

                # Step 3: Reconstruct the function on the sphere

        # Create a spherical_array to hold the supertranslated waveform
        supertransl_spherical_waveform = spherical_array(grid_info=grid_info)

        # Set the data.
        supertransl_spherical_waveform.data = supertransl_spherical_grid

        # Get modes_array from spherical_array
        Psi4_supertransl_modes = supertransl_spherical_waveform.to_modes_array(
            ell_max=ell_max
        )

        return Psi4_supertransl_modes

    def boost(self, conformal_factor, grid_info=None):
        """Boost the waveform given the unboosted waveform
        and the boost conformal factor.

        Parameters
        ----------
        conformal_factor: 2d array
                           The conformal factor for
                           the Lorentz transformation.
                           It may be a single floating point number
                           on a spherical grid. The array will be
                           of dimensions [ntheta, nphi].

        Returns
        -------
        boosted_waveform: spherical_array
                           The class instance `spherical_array`
                           that contains the boosted waveform.
        """

        from spectral.spherical.grids import UniformGrid

        # Construct a spherical grid.
        if grid_info is None:
            grid_info = UniformGrid()

        # Get spherical array from modes.
        unboosted_waveform = self.to_spherical_array(grid_info)

        boosted_waveform_data = unboosted_waveform.boost(conformal_factor)

        # Construct a 2d waveform array object
        boosted_waveform = spherical_array(
            grid_info=unboosted_waveform.grid_info,
            data=np.array(boosted_waveform_data),
        )
        boosted_waveform.label = "boosted"

        # Get modes from spherical data.
        # boosted_waveform_modes = boosted_waveform.to_modes_array()

        # return boosted_waveform_modes
        return boosted_waveform

    def get_strain_from_psi4(self, omega0="auto"):
        """Get the strain `modes_array` from :math:`\\Psi_4` by
        fixed frequency integration.

        Parameters
        ----------
        omega0: float, optional
                The lower cutoff angular frequency for FFI.
                By default, it computes this from the mode data.

        Returns
        -------
        strain_waveform: modes_array
                          The computed strain modes.
        """

        # Initialize a mode array for strain.
        strain_waveform = modes_array(
            label="{} strain from Psi4".format(self.label),
            r_ext=self.r_ext,
            ell_max=8,
            modes_list=self.modes_list,
        )

        strain_waveform.ell_max = self.ell_max

        data_len = self.data_len

        strain_waveform.create_modes_array(
            ell_max=self.ell_max, data_len=data_len
        )

        # Integrate,
        from waveformtools.integrate import fixed_frequency_integrator
        from waveformtools.waveformtools import (
            get_starting_angular_frequency as sang_f,
        )

        omega_st = omega0
        for item in self.modes_list:
            ell, emm_list = item
            for emm in emm_list:
                mode_data = self.mode(ell, emm)
                if omega0 == "auto":
                    omega_st = (
                        abs(sang_f(mode_data, delta_t=self.delta_t())) / 10
                    )
                strain_time, strain_mode_data = fixed_frequency_integrator(
                    udata_time=mode_data,
                    delta_t=self.delta_t(),
                    omega0=omega_st,
                    order=2,
                )
                strain_waveform.set_mode_data(ell, emm, data=strain_mode_data)

        strain_waveform.time_axis = strain_time

        strain_waveform.trim(trim_upto_time=0)

        return strain_waveform

    def get_news_from_psi4(self, omega0="auto"):
        """Get the News `modes_array` from :math:`\\Psi_4` by
        fixed frequency integration.

        Parameters
        ----------
        omega0: float, optional
                 The lower cutoff angular frequency for FFI.
                 By default, it computes this as one tenth of
                 the starting frequency of the respective mode data.

        Returns
        -------
        news_waveform: modes_array
                        The computed strain modes.
        """

        # Initialize a mode array for strain.
        # news_waveform = modes_array(label=f'{self.label} news from
        # Psi4', r_ext=500, ell_max=8, modes_list=self.modes_list)
        news_waveform = modes_array(
            label="{} news from Psi4".format(self.label),
            r_ext=500,
            ell_max=8,
            modes_list=self.modes_list,
        )

        # news_waveform.time_axis = self.time_axis
        news_waveform.ell_max = self.ell_max

        data_len = self.data_len

        news_waveform.create_modes_array(
            ell_max=self.ell_max, data_len=data_len
        )

        # Integrate,
        omega_st = omega0
        for item in self.modes_list:
            ell, emm_list = item
            for emm in emm_list:
                mode_data = self.mode(ell, emm)
                if omega0 == "auto":
                    omega_st = (
                        abs(sang_f(mode_data, delta_t=self.delta_t())) / 10
                    )
                news_time_axis, news_mode_data = fixed_frequency_integrator(
                    udata_time=mode_data,
                    delta_t=self.delta_t(),
                    omega0=omega_st,
                    order=1,
                )
                news_waveform.set_mode_data(
                    ell_value=ell, emm_value=emm, data=news_mode_data
                )

        news_waveform._time_axis = news_time_axis

        news_waveform.trim(trim_upto_time=0)

        return news_waveform

    def taper(self, zeros="auto"):
        """Taper a waveform at both ends and insert zeros if needed

        Parameters
        ----------
        zeros: int
                The number of zeros to add at rach end

        Returns
        -------
        tapered_modes: modes_array
                        Modes array with tapered mode data.
        """

        from waveformtools.waveformtools import taper

        if zeros == "auto":
            # Decide the number of zeros
            data_len = self.data_len

            nearest_power = int(np.log(data_len) / np.log(2))
            req_len = np.power(2, nearest_power + 1)
            zeros = req_len - data_len
            message("num_zeros", zeros)

        # New modes array.

        tapered_modes = None

        for item in self.modes_list[:]:
            ell, emm_list = item
            for emm in emm_list:
                input_data_re = self.mode(ell, emm).real
                input_data_im = self.mode(ell, emm).imag

                tapered_data_re = taper(input_data_re, zeros=zeros)
                tapered_data_im = taper(input_data_im, zeros=zeros)

                # tapered_data_re = taper_tanh(input_data_re,
                # delta_t=self.delta_t())
                # tapered_data_im = taper_tanh(input_data_im,
                # delta_t=self.delta_t())

                new_data_len = len(tapered_data_re)

                if tapered_modes is None:
                    tapered_modes = modes_array(
                        label="tapered {}".format(self.label),
                        r_ext=self.r_ext,
                        modes_list=self.modes_list,
                        ell_max=self.ell_max,
                    )

                    tapered_modes.create_modes_array(
                        ell_max=self.ell_max, data_len=new_data_len
                    )
                tapered_data = tapered_data_re + 1j * tapered_data_im

                # message(len(tapered_data_re))
                tapered_modes.set_mode_data(ell, emm, data=tapered_data)

        # Set the time axis
        new_time_axis = np.arange(
            0, new_data_len * self.delta_t(), self.delta_t()
        )

        tapered_modes.time_axis = new_time_axis

        # Recenter the modes.
        tapered_modes.trim(trim_upto_time=0)

        return tapered_modes

    def taper_tanh(
        self, time_axis=None, zeros="auto", duration=10, sides="both"
    ):
        """Taper a waveform at both ends and insert zeros if needed

        Parameters
        ----------
        zeros: int
                The number of zeros to add at rach end

        Returns
        -------
        tapered_modes: modes_array
                        Modes array with tapered mode data.
        """

        from waveformtools.waveformtools import taper_tanh

        if zeros == "auto":
            # Decide the number of zeros
            data_len = self.data_len

            nearest_power = int(np.log(data_len) / np.log(2))
            req_len = np.power(2, nearest_power + 1)
            zeros = req_len - data_len
            # message('num_zeros', zeros)

        # New modes array.

        tapered_modes = None

        for item in self.modes_list[:]:
            ell, emm_list = item
            for emm in emm_list:
                input_data_re = self.mode(ell, emm).real
                input_data_im = self.mode(ell, emm).imag

                # tapered_data_re = taper(input_data_re, zeros=zeros)
                # tapered_data_im = taper(input_data_im, zeros=zeros)

                _, tapered_data_re = taper_tanh(
                    input_data_re,
                    delta_t=self.delta_t(),
                    duration=duration,
                    sides=sides,
                )
                _, tapered_data_im = taper_tanh(
                    input_data_im,
                    delta_t=self.delta_t(),
                    duration=duration,
                    sides=sides,
                )

                new_data_len = len(tapered_data_re)

                if tapered_modes is None:
                    tapered_modes = modes_array(
                        label="tapered {}".format(self.label),
                        r_ext=self.r_ext,
                        modes_list=self.modes_list,
                        ell_max=self.ell_max,
                    )

                    tapered_modes.create_modes_array(
                        ell_max=self.ell_max, data_len=new_data_len
                    )
                tapered_data = tapered_data_re + 1j * tapered_data_im

                # message(len(tapered_data_re))
                tapered_modes.set_mode_data(ell, emm, data=tapered_data)

        # Set the time axis
        new_time_axis = np.arange(
            0, new_data_len * self.delta_t(), self.delta_t()
        )

        tapered_modes.time_axis = new_time_axis

        # Recenter the modes.
        tapered_modes.trim(trim_upto_time=0)

        return tapered_modes

    def low_cut(self, omega0=0.03, order=2):
        """Apply the low cut filter from `waveformtools.low_cut_filter`

        Parameters
        ----------
        order: int, optional
                The order of the butterworth filter.
        omega0: float, optional
                 The cutoff frequency of the butterworth filter.

        Returns:
        --------
        filtered_modes: modes_array
                         A modes array object containing filtered modes.
        """

        # modes_array for filtered data.
        filtered_modes = None

        # Import the filter
        from waveformtools.waveformtools import low_cut_filter

        for item in self.modes_list:
            # Iterate over available modes.
            ell, emm_list = item
            for emm in emm_list:
                if filtered_modes is None:
                    # Create filtered_modes
                    filtered_modes = modes_array(
                        label="lc filtered {}".format(self.label),
                        r_ext=self.r_ext,
                        modes_list=self.modes_list,
                        ell_max=self.ell_max,
                    )

                    filtered_modes.create_modes_array(
                        ell_max=self.ell_max, data_len=self.data_len
                    )

                # Get filtered mode data.
                filtered_data = low_cut_filter(
                    self.mode(ell, emm),
                    self.frequency_axis,
                    omega0=omega0,
                    order=order,
                )

                # Set the mode data.
                filtered_modes.set_mode_data(ell, emm, data=filtered_data)

        # Set the f axis.
        filtered_modes.frequency_axis = self.frequency_axis

        return filtered_modes

    def to_td_waveform(
        self,
        Mtotal=1,
        theta=0,
        phi=0,
        alpha=None,
        distance=1,
        delta_t=None,
        method="precise",
        k=None,
    ):
        """Get the plus and cross polarizations of
        of the waveform time series by summing the modes.

        Parameters
        ----------
        theta, phi: float
                    The inclination and the azimuthal
                    angular position of the observer
                    in the NR coodinate system.
        distance: float
                   The distance to the source

        method: str
                 The method to use to generate
                 the SWSH basis. This can be
                 `precise` or `fast`.

        Returns
        -------
        taxis, hp, hc: 1d arrays
                        The 1d arrays of the time axis and
                        the polarizations of the waveforms.

        Notes
        -----
        This does not rotate the polarizations. The rotation
        must be done separately using the `rotate_polarizations`
        function of `waveformtools.transforms`

        For precessing systems and to obtain the waveform
        in the LAL convention, one should use the
        nrcatalogtools package to obtain the correct
        angles first.
        """

        if method == "fast":
            message("Using fast SWSH method")
            from spectral.spherical.swsh import Yslm_vec as Yslm
        elif method == "precise":
            message("Using precise SWSH method")
            from spectral.spherical.swsh import Yslm_prec as Yslm
        else:
            raise NotImplementedError(f"Unknown method {method}")

        wts = 0
        all_yslm = []
        flat_modes = []

        for ell, emm_list in self.modes_list:
            for emm in emm_list:
                Y = Yslm(self.spin_weight, ell, emm, theta=theta, phi=phi)

                all_yslm.append(np.complex256(Y))
                flat_modes.append(self.mode(ell, emm))

        all_yslm = np.complex256(all_yslm)

        wts = np.dot(all_yslm, flat_modes)

        # Rescale the time axis
        taxis = self.time_axis * tuc * Mtotal

        if isinstance(delta_t, float):
            # Resample the waveform
            new_taxis = np.arange(taxis[0], taxis[-1], delta_t)
            wts = interp_resam_wfs(wts, taxis, new_taxis, k=k)
            taxis = new_taxis

        # Rescale the magnitude of the waveform
        wts = Mtotal * wts / (distance * dMpc * muc)

        # Rotate polarizations
        if alpha is not None:
            from waveformtools.transforms import rotate_polarizations

            wts = rotate_polarizations(wts, alpha)

        return taxis, wts.real, -wts.imag

    def evaluate_angular(
        self, theta=None, phi=None, ell_max=None, max_t_steps=None
    ):
        """Evaluate the expansion at requested angular coordinates
        by generating SWSHs in parallel and vectorizing the
        summation"""
        from spectral.spherical.Yslm_mp import Yslm_mp

        if ell_max is None:
            ell_max = self.ell_max

        if (np.array(theta) == np.array(None)).all() or (
            np.array(phi) == np.array(None)
        ).all():
            theta, phi = self.Grid.meshgrid

        sYlm = Yslm_mp(
            ell_max=ell_max, spin_weight=self.spin_weight, theta=theta, phi=phi
        )
        sYlm.run()
        Ylm_vec = sYlm.sYlm_modes._modes_data
        modes_data_len = len(Ylm_vec)

        val = np.tensordot(
            self._modes_data[:modes_data_len, ..., :max_t_steps],
            Ylm_vec,
            axes=((0), (0)),
        )

        return val

    def plot_modes(self, modes_to_plot=[(2, 2), (3, 3)]):
        """Plot the requested set of modes"""

        import matplotlib.pyplot as plt

        try:
            import config

            config.conf_matplolib()
        except Exception as excep:
            excep("Unable to conf matplotlib")

        fig, ax = plt.subplots(figsize=(12, 6))

        for one_mode in modes_to_plot:
            ell, emm = one_mode

            ax.plot(
                self.time_axis,
                self.mode(ell, emm).real,
                label=f"Re( ({ell}, {emm}) )",
            )

        ax.set_xlabel("t/M")
        ax.set_ylabel(r"$\Psi_{(\ell, m)}$")
        plt.legend()
        plt.show()
