# Copyright 2023 Infleqtion
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
from scipy.interpolate import RectBivariateSpline, interp1d

from oqtant.simulator.three_dim import ThreeDim


class WaveFunction:
    """
    'WaveFunction' Defines representation for a wavefunction
    """

    def __init__(self):
        self.tdm = ThreeDim()
        initial_psi = self.initial_psi().astype(complex)
        self.psi = self.normalize(initial_psi)

    def initial_psi(self, sigma_x=1, sigma_r=1) -> np.ndarray:
        """
        Defines the initial wave function of the system with
        controllable widths.

        Args:
            sigma_x int:  width in the x direction
            sigma_r int:  width in the r direction

        Returns:
            psi np.ndarray:  initial wave function

        """
        init_x = np.exp(-((self.tdm.x) ** 2) / (2 * sigma_x**2))
        init_r = np.exp(-((self.tdm.r) ** 2) / (2 * sigma_r**2))

        return init_x * init_r

    def normalize(self, psi):
        """
        Normalizes the wave function to the number of atoms
        Applies normalization directly to self.psi

        Args:
            psi np.ndarray:  wavefunction

        Returns:
            None
        """
        B = (
            2
            * np.pi
            * self.tdm.dx
            * self.tdm.dr
            * np.sum(self.tdm.r * np.abs(psi) ** 2)
        )
        # In these units, psi is normalized to unity.  U0 is then the TF parameter N*a/aho.
        return psi * np.sqrt(1 / B)

    @property
    def density(self):
        """
        Returns the density of the wave function
        """
        return np.abs(self.psi) ** 2

    @property
    def column_densities(self):
        """
        Returns the column densities (#/Length^2)
        column_zy:  the zy axes
        column_zx:  the zx axes
        profiles:  the profiles in the x and z axes
        """
        # column_xz
        column_zy = np.zeros((2 * self.tdm.Nr, 2 * self.tdm.Nr))
        profile_r = self.tdm.dx * np.sum(np.abs(self.psi) ** 2, axis=1)  #
        # 1D interpolation
        profile_r_interp = interp1d(
            self.tdm.r[:, 0],
            profile_r,
            fill_value=0.0,
            bounds_error=False,
            kind="cubic",
        )

        for iz in range(len(self.tdm.z[0, :])):
            for iy in range(len(self.tdm.y[:, 0])):
                r = np.sqrt(self.tdm.z[0, iz] ** 2 + self.tdm.y[iy, 0] ** 2)
                column_zy[iz, iy] = profile_r_interp(r)

        # column zy
        # 2D interpolation
        density_interp = RectBivariateSpline(
            self.tdm.r[:, 0], self.tdm.x[0, :], self.density, kx=3, ky=3
        )  # cubic spline interpolation
        column_zx = np.zeros((2 * self.tdm.Nr, self.tdm.Nx))

        for iz in range(len(self.tdm.z[0, :])):
            for ix in range(len(self.tdm.x[0, :])):
                r = np.sqrt(self.tdm.z[0, iz] ** 2 + self.tdm.y[:, 0] ** 2)
                density_slice = density_interp(np.sort(r), self.tdm.x[0, ix])
                column_zx[iz, ix] = self.tdm.dy * np.sum(density_slice)

        #
        profile_z_zy = self.tdm.dy * np.sum(column_zy, axis=1)
        profile_z_zx = self.tdm.dx * np.sum(column_zx, axis=1)
        profile_x = self.tdm.dz * np.sum(column_zx, axis=0)
        profiles = [profile_z_zy, profile_z_zx, profile_x]
        # check that interpolation preserved norm
        if (
            self.tdm.dy * sum(profile_z_zy) > 0.95
            and self.tdm.dx * sum(profile_x) > 0.95
        ):
            pass
        else:
            raise ValueError("Lost norm in interpolation")

        # the profile_x should match up with the output of the profile function.
        return (
            column_zy,
            column_zx,
            profiles,
        )

    @property
    def density_profiles(self):
        """
        Returns the density profiles along the x and r axes.  #/Length
        These match the integrated column densities
        """
        profile_x = (
            2 * np.pi * self.tdm.dr * np.sum(self.tdm.r * np.abs(self.psi) ** 2, axis=0)
        )
        profile_r = (
            2
            * np.pi
            * self.tdm.dx
            * self.tdm.r[:, 0]
            * np.sum(np.abs(self.psi) ** 2, axis=1)
        )
        return profile_x, profile_r

    @property
    def atom_number(self):
        """
        Returns atom number
        """
        return (
            2
            * np.pi
            * self.tdm.dx
            * self.tdm.dr
            * np.sum(self.tdm.r * np.abs(self.psi) ** 2)
        )

    @property
    def phase(self):
        """
        Returns the phase of the wave function
        """
        return np.angle(self.psi)

    @property
    def flow(self):
        """
        Returns the superfluid velocities in X and R directions
        Returns two 2D arrays
        """
        diff_r = (
            2 * self.tdm.dr * np.gradient(self.phase, self.tdm.dr, axis=0, edge_order=1)
        )  # axis 0 for r
        diff_x = (
            2 * self.tdm.dz * np.gradient(self.phase, self.tdm.dz, axis=1, edge_order=1)
        )  # axis 1 for x

        # dealing with 'jumps' from negative pi to pi.  Essentially enforcing that diff_i resides between -pi and pi
        diff_r[diff_r > np.pi] -= 2 * np.pi
        diff_r[diff_r < -np.pi] += 2 * np.pi
        diff_x[diff_x > np.pi] -= 2 * np.pi
        diff_x[diff_x < -np.pi] += 2 * np.pi

        flow_r = 2 * np.pi * diff_r / (2 * self.tdm.dr)
        flow_x = 2 * np.pi * diff_x / (2 * self.tdm.dx)

        return flow_r, flow_x

    @property
    def current(self):
        """
        Returns the total current along the X-direction
        Returns a 1D array
        """
        _, flow_x = self.flow  # v_z(r,x)
        current_x = (
            2 * np.pi * self.tdm.dr * np.sum(self.tdm.r * flow_x * self.density, axis=0)
        )  # j_z(z):  A 1D array of current along the x-direction.
        return current_x

    @property
    def com_position(self):
        """
        Returns the center of mass coordinates of the cloud in X direction
        The center of mass cannot be displaced in the radial direction by assumption
        Returns 1 scalar
        Useful for diagnostics
        """
        com_x = (
            2
            * np.pi
            * self.tdm.dx
            * self.tdm.dr
            * np.sum(self.tdm.x * self.tdm.r * np.abs(self.psi) ** 2)
        )

        return com_x

    @property
    def widths(self):
        """
        Returns the widths (\\Delta r and \\Delta x) in radial and x directions
        Returns 2 scalars
        """
        square_width_x = (
            2
            * np.pi
            * self.tdm.dr
            * self.tdm.dx
            * np.sum(self.tdm.r * self.tdm.x**2 * np.abs(self.psi) ** 2)
        )
        square_width_r = (
            2
            * np.pi
            * self.tdm.dr
            * self.tdm.dx
            * np.sum(self.tdm.r**3 * np.abs(self.psi) ** 2)
        )
        com_x = self.com_position()

        return np.sqrt(square_width_x - com_x**2), np.sqrt(square_width_r)
