import aquila_borg as borg
import jax
import jax.numpy as jnp
from functools import partial

from borg_velocity.utils import myprint

class NullForward(borg.forward.BaseForwardModel):
    """
    BORG forward model which does nothing but stores
    the values of parameters to be used by the likelihood
    """
    def __init__(self, box: borg.forward.BoxModel) -> None:
        """
        Initialise the NullForward class
        Args:
            box (borg.forward.BoxModel): The input box model.
        """
        super().__init__(box, box)
        self.setName("nullforward")

        self.params = {}
        self.setCosmoParams(borg.cosmo.CosmologicalParameters())
        cosmo = self.getCosmoParams()
        cosmo.n_s = 0.96241
        self.setCosmoParams(cosmo)


    def setModelParams(self, params: dict) -> None:
        """
        Change the values of the model parameters to those given by params

        Args:
            params (dict): Dictionary of updated model parameters.
        """
        for k, v in params.items():
            self.params[k] = v
        myprint(f'Updated model parameters: {self.params}')

    def getModelParam(self, model, keyname: str):
        """
        This queries the current state of the parameters keyname in model model.

        Args:
            model: The model
            keyname (str): The name of the parameter of interest
        """
        return self.params[keyname]


@partial(jax.jit, static_argnums=(2,3))
def dens2vel_linear(delta, f, Lbox, smooth_R):
    """
    Converts the overdensity field, delta, into
    the velocity field, v, in units of [???], assuming
    the linear relation between the two.
    
    Args:
        :delta (jnp.ndarray): (N,N,N) array of overdensity field
        :f (float): The growth rate
        :Lbox (float): The box length in Mpc/h
        :smooth_R (float): The smoothing scale in Mpc/h
        
    Returns:
        :v (jnp.ndarray): The velocity field defined on the grid
        
    TO CHECK: UNITS
    QUESTION: Why two different filters?
    """
    
    N = delta.shape[0]
    V = Lbox**3
    dV = (Lbox / N) ** 3
    
    # Forward fft
    delta_k = dV / V * jnp.fft.rfftn(delta)
    delta_k = jnp.array([delta_k.real, delta_k.imag])
    
    # Fourier mask for Nyquist planes
    N_Z = N//2 + 1
    Fourier_mask = jnp.ones((N,N,N_Z))
    Fourier_mask = Fourier_mask.at[:,N_Z,0].set(0.)
    Fourier_mask = Fourier_mask.at[:,N_Z,-1].set(0.)
    
    # Indices for symmetrising
    update_index_real_0  = jnp.index_exp[0,N_Z:,:,0]
    update_index_real_ny = jnp.index_exp[0,N_Z:,:,N_Z-1]
    update_index_imag_0  = jnp.index_exp[1,N_Z:,:,0]
    update_index_imag_ny = jnp.index_exp[1,N_Z:,:,N_Z-1]
    flip_indices = -jnp.arange(N)
    flip_indices = flip_indices.at[N_Z-1].set(-flip_indices[N_Z-1])
    # flip_indices = jnp.array(flip_indices.tolist())
    
    # Symmetrise
    delta_k = Fourier_mask[jnp.newaxis] * delta_k
    delta_k = delta_k.at[update_index_real_0].set(
            jnp.take(jnp.flip(delta_k[0, 1:(N_Z-1), :, 0], axis=0), 
                     flip_indices, axis=1)
        )
    delta_k = delta_k.at[update_index_real_ny].set(
            jnp.take(jnp.flip(delta_k[0,1:(N_Z-1),:,N_Z-1], axis=0), 
                     flip_indices, axis=1)
        )
    delta_k = delta_k.at[update_index_imag_0].set(
        -jnp.take(jnp.flip(delta_k[1,1:(N_Z-1),:,0], axis=0), 
                  flip_indices, axis=1)
        )
    
    delta_k = delta_k[0] + jnp.array(complex(0, 1)) * delta_k[1]
    
    # Get k grid
    k = 2*jnp.pi*jnp.fft.fftfreq(N, d=Lbox/N)
    k_norm = jnp.sqrt(k[:,None,None]**2 + k[None,:,None]**2 + k[None,None,:N_Z]**2)
    k_norm = jnp.where(k_norm < 1e-10, 1e-15, k_norm)
    
    # Filter
    k_filter = jnp.exp(-0.5 * (k_norm[:,:,:N_Z] * smooth_R) ** 2)
    smooth_filter = k_filter
    vx = (
        smooth_filter * k_filter
        * jnp.array(complex(0, 1)) * 100 * f
        * delta_k
        * jnp.tile(k[:,None,None],(1,N,N_Z))
        / k_norm**2
    )
    vx = (jnp.fft.irfftn(vx)*V/dV)
    
    vy = (
        smooth_filter * k_filter
        * jnp.array(complex(0, 1)) * 100 * f
        * delta_k
        * jnp.tile(k[None,:,None], (N,1,N_Z))
        / k_norm**2
    )
    vy = (jnp.fft.irfftn(vy)*V/dV)
    
    vz = (
        smooth_filter * k_filter
        * jnp.array(complex(0, 1)) * 100 * f
        * delta_k
        * jnp.tile(k[None,None,:N_Z], (N,N,1))
        / k_norm**2
    )
    vz = (jnp.fft.irfftn(vz)*V/dV)
    
    return jnp.array([vx, vy, vz])