"""
Computations of the eigenvalues and eigenvectors for the Euler (hydrodynamics) equations.

Problems for differentiation largely follow from square roots and divisions:
The derivative of sqrt(x) is 1/(2*sqrt(x)) and of 1/x is -1/x^2, where both expressions
are problematic for small x, especially when multiplying gradients in the backward pass,
-> exploding gradients.
"""

from functools import partial
import jax
import jax.numpy as jnp
from typing import Union

from astronomix._stencil_operations._stencil_operations import _shift
from astronomix.variable_registry.registered_variables import RegisteredVariables


def diff_safe_sqrt(x):
    epsilon = 1e-12
    x_safe = jnp.maximum(x, epsilon)
    return jnp.sqrt(x_safe)


@partial(jax.jit, static_argnames=["registered_variables"])
def _eigenvalue_building_blocks(
    conserved_state,
    gamma,
    rhomin,
    pgmin,
    registered_variables: RegisteredVariables,
):
    # unpack the conserved variables
    density = conserved_state[registered_variables.density_index]
    momentum_x = conserved_state[registered_variables.momentum_index.x]
    momentum_y = conserved_state[registered_variables.momentum_index.y]
    momentum_z = conserved_state[registered_variables.momentum_index.z]
    energy = conserved_state[registered_variables.energy_index]

    # compute primitives
    rho = density
    velocity_x = momentum_x / rho
    velocity_y = momentum_y / rho
    velocity_z = momentum_z / rho
    velocity_squared = (
        velocity_x * velocity_x + velocity_y * velocity_y + velocity_z * velocity_z
    )
    
    gas_pressure = (gamma - 1.0) * (
        energy - 0.5 * rho * velocity_squared
    )

    # redefine the density and pressure, and energy based on floors
    rho = jnp.where(
        (rho < rhomin) | (gas_pressure < pgmin), jnp.maximum(rho, rhomin), rho
    )
    gas_pressure = jnp.where(
        (rho < rhomin) | (gas_pressure < pgmin),
        jnp.maximum(gas_pressure, pgmin),
        gas_pressure,
    )
    energy = jnp.where(
        (rho < rhomin) | (gas_pressure < pgmin),
        gas_pressure / (gamma - 1.0) + 0.5 * rho * velocity_squared,
        energy,
    )

    # compute derived quantities
    sound_speed = diff_safe_sqrt(jnp.maximum(0.0, gamma * jnp.abs(gas_pressure / rho)))

    return (
        velocity_x,
        sound_speed,
    )

@partial(jax.jit, static_argnames=["registered_variables"])
def _eigenvector_building_blocks(
    conserved_state,
    gamma,
    rhomin,
    pgmin,
    registered_variables: RegisteredVariables,
):
    # unpack conserved variables
    rho = conserved_state[registered_variables.density_index]  
    momentum_x = conserved_state[registered_variables.momentum_index.x]
    momentum_y = conserved_state[registered_variables.momentum_index.y]
    momentum_z = conserved_state[registered_variables.momentum_index.z]
    energy = conserved_state[registered_variables.energy_index]

    # compute primitives
    velocity_x = momentum_x / rho
    velocity_y = momentum_y / rho
    velocity_z = momentum_z / rho
    velocity_sq = (
        velocity_x * velocity_x + velocity_y * velocity_y + velocity_z * velocity_z
    )
    
    gas_pressure = (gamma - 1.0) * (energy - 0.5 * rho * velocity_sq)

    # redefine the density and pressure, and energy based on floors
    rho = jnp.where(
        (rho < rhomin) | (gas_pressure < pgmin), jnp.maximum(rho, rhomin), rho
    )
    gas_pressure = jnp.where(
        (rho < rhomin) | (gas_pressure < pgmin),
        jnp.maximum(gas_pressure, pgmin),
        gas_pressure,
    )
    energy = jnp.where(
        (rho < rhomin) | (gas_pressure < pgmin),
        gas_pressure / (gamma - 1.0) + 0.5 * rho * velocity_sq,
        energy,
    )

    specific_enthalpy = (energy + gas_pressure) / rho

    # periodic average to interfaces
    def avg_x(arr):
        return 0.5 * (arr + _shift(arr, shift=-1, axis=0))

    rho_interface = avg_x(rho)
    velocity_x_interface = avg_x(velocity_x)
    velocity_y_interface = avg_x(velocity_y)
    velocity_z_interface = avg_x(velocity_z)
    specific_enthalpy_interface = avg_x(specific_enthalpy)

    # interface derived quantities
    velocity_sq_interface = (
        velocity_x_interface * velocity_x_interface
        + velocity_y_interface * velocity_y_interface
        + velocity_z_interface * velocity_z_interface
    )

    # enthalpy based sound speed at interfaces
    sound_speed_sq_interface = (gamma - 1.0) * (
        specific_enthalpy_interface - 0.5 * velocity_sq_interface
    )
    sound_speed_interface = diff_safe_sqrt(jnp.maximum(0.0, sound_speed_sq_interface))

    sound_speed_sq_inverse = jnp.where(
        sound_speed_sq_interface > 0.0, 1.0 / sound_speed_sq_interface, 0.0
    )

    gamma_minus_one = gamma - 1.0

    return (
        rho_interface,
        velocity_x_interface,
        velocity_y_interface,
        velocity_z_interface,
        velocity_sq_interface,
        specific_enthalpy_interface,
        sound_speed_interface,
        sound_speed_sq_interface,
        sound_speed_sq_inverse,
        gamma_minus_one,
    )


@partial(jax.jit, static_argnames=["registered_variables"])
def _eigen_R_col_hydro(
    conserved_state,
    rhomin: Union[float, jnp.ndarray],
    pgmin: Union[float, jnp.ndarray],
    gamma: Union[float, jnp.ndarray],
    registered_variables: RegisteredVariables,
    col: int,
):
    (
        rho_interface,
        velocity_x_interface,
        velocity_y_interface,
        velocity_z_interface,
        velocity_sq_interface,
        specific_enthalpy_interface,
        sound_speed_interface,
        sound_speed_sq_interface,
        sound_speed_sq_inverse,
        gamma_minus_one,
    ) = _eigenvector_building_blocks(
        conserved_state,
        gamma,
        rhomin,
        pgmin,
        registered_variables,
    )

    # shorter names for registry indices
    density_index = registered_variables.density_index
    momentum_index_x = registered_variables.momentum_index.x
    momentum_index_y = registered_variables.momentum_index.y
    momentum_index_z = registered_variables.momentum_index.z
    energy_index = registered_variables.energy_index

    def col_0():
        # Column 1 (Left Acoustic / u - c)
        R = jnp.zeros_like(conserved_state)
        R = R.at[density_index].set(1.0)
        R = R.at[momentum_index_x].set(velocity_x_interface - sound_speed_interface)
        R = R.at[momentum_index_y].set(velocity_y_interface)
        R = R.at[momentum_index_z].set(velocity_z_interface)
        R = R.at[energy_index].set(specific_enthalpy_interface - velocity_x_interface * sound_speed_interface)
        return R

    def col_1():
        # Column 2 (Entropy / u)
        R = jnp.zeros_like(conserved_state)
        R = R.at[density_index].set(1.0)
        R = R.at[momentum_index_x].set(velocity_x_interface)
        R = R.at[momentum_index_y].set(velocity_y_interface)
        R = R.at[momentum_index_z].set(velocity_z_interface)
        R = R.at[energy_index].set(0.5 * velocity_sq_interface)
        return R

    def col_2():
        # Column 3 (Shear Y / u)
        R = jnp.zeros_like(conserved_state)
        R = R.at[density_index].set(0.0)
        R = R.at[momentum_index_x].set(0.0)
        R = R.at[momentum_index_y].set(1.0)
        R = R.at[momentum_index_z].set(0.0)
        R = R.at[energy_index].set(velocity_y_interface)
        return R

    def col_3():
        # Column 4 (Shear Z / u)
        R = jnp.zeros_like(conserved_state)
        R = R.at[density_index].set(0.0)
        R = R.at[momentum_index_x].set(0.0)
        R = R.at[momentum_index_y].set(0.0)
        R = R.at[momentum_index_z].set(1.0)
        R = R.at[energy_index].set(velocity_z_interface)
        return R

    def col_4():
        # Column 5 (Right Acoustic / u + c)
        R = jnp.zeros_like(conserved_state)
        R = R.at[density_index].set(1.0)
        R = R.at[momentum_index_x].set(velocity_x_interface + sound_speed_interface)
        R = R.at[momentum_index_y].set(velocity_y_interface)
        R = R.at[momentum_index_z].set(velocity_z_interface)
        R = R.at[energy_index].set(specific_enthalpy_interface + velocity_x_interface * sound_speed_interface)
        return R

    R = jax.lax.switch(col, [col_0, col_1, col_2, col_3, col_4])

    return R


@partial(jax.jit, static_argnames=["registered_variables"])
def _eigen_L_row_hydro(
    conserved_state,
    rhomin: Union[float, jnp.ndarray],
    pgmin: Union[float, jnp.ndarray],
    gamma: Union[float, jnp.ndarray],
    registered_variables: RegisteredVariables,
    row: int,
):
    (
        rho_interface,
        velocity_x_interface,
        velocity_y_interface,
        velocity_z_interface,
        velocity_sq_interface,
        specific_enthalpy_interface,
        sound_speed_interface,
        sound_speed_sq_interface,
        sound_speed_sq_inverse,
        gamma_minus_one,
    ) = _eigenvector_building_blocks(
        conserved_state,
        gamma,
        rhomin,
        pgmin,
        registered_variables,
    )

    # shorter names for registry indices
    density_index = registered_variables.density_index
    momentum_index_x = registered_variables.momentum_index.x
    momentum_index_y = registered_variables.momentum_index.y
    momentum_index_z = registered_variables.momentum_index.z
    energy_index = registered_variables.energy_index

    def row_0():
        # Row 1 (Left Acoustic / u - c)
        L = jnp.zeros_like(conserved_state)
        L = L.at[density_index].set(0.5 * gamma_minus_one * velocity_sq_interface + velocity_x_interface * sound_speed_interface)
        L = L.at[momentum_index_x].set(-(gamma_minus_one * velocity_x_interface + sound_speed_interface))
        L = L.at[momentum_index_y].set(-gamma_minus_one * velocity_y_interface)
        L = L.at[momentum_index_z].set(-gamma_minus_one * velocity_z_interface)
        L = L.at[energy_index].set(gamma_minus_one)
        L = 0.5 * sound_speed_sq_inverse * L
        return L

    def row_1():
        # Row 2 (Entropy / u)
        L = jnp.zeros_like(conserved_state)
        L = L.at[density_index].set(sound_speed_sq_interface - 0.5 * gamma_minus_one * velocity_sq_interface)
        L = L.at[momentum_index_x].set(gamma_minus_one * velocity_x_interface)
        L = L.at[momentum_index_y].set(gamma_minus_one * velocity_y_interface)
        L = L.at[momentum_index_z].set(gamma_minus_one * velocity_z_interface)
        L = L.at[energy_index].set(-gamma_minus_one)
        L = sound_speed_sq_inverse * L
        return L

    def row_2():
        # Row 3 (Shear Y / u)
        L = jnp.zeros_like(conserved_state)
        L = L.at[density_index].set(-velocity_y_interface)
        L = L.at[momentum_index_x].set(0.0)
        L = L.at[momentum_index_y].set(1.0)
        L = L.at[momentum_index_z].set(0.0)
        L = L.at[energy_index].set(0.0)
        return L

    def row_3():
        # Row 4 (Shear Z / u)
        L = jnp.zeros_like(conserved_state)
        L = L.at[density_index].set(-velocity_z_interface)
        L = L.at[momentum_index_x].set(0.0)
        L = L.at[momentum_index_y].set(0.0)
        L = L.at[momentum_index_z].set(1.0)
        L = L.at[energy_index].set(0.0)
        return L

    def row_4():
        # Row 5 (Right Acoustic / u + c)
        L = jnp.zeros_like(conserved_state)
        L = L.at[density_index].set(0.5 * gamma_minus_one * velocity_sq_interface - velocity_x_interface * sound_speed_interface)
        L = L.at[momentum_index_x].set(-(gamma_minus_one * velocity_x_interface - sound_speed_interface))
        L = L.at[momentum_index_y].set(-gamma_minus_one * velocity_y_interface)
        L = L.at[momentum_index_z].set(-gamma_minus_one * velocity_z_interface)
        L = L.at[energy_index].set(gamma_minus_one)
        L = 0.5 * sound_speed_sq_inverse * L
        return L

    L = jax.lax.switch(row, [row_0, row_1, row_2, row_3, row_4])

    return L


@partial(jax.jit, static_argnames=["registered_variables"])
def _eigen_all_lambdas_hydro(
    conserved_state,
    rhomin: Union[float, jnp.ndarray],
    pgmin: Union[float, jnp.ndarray],
    gamma: Union[float, jnp.ndarray],
    registered_variables: RegisteredVariables,
):
    (
        velocity_x,
        sound_speed,
    ) = _eigenvalue_building_blocks(
        conserved_state,
        gamma,
        rhomin,
        pgmin,
        registered_variables,
    )

    return jnp.stack(
        [
            velocity_x - sound_speed,
            velocity_x,
            velocity_x,
            velocity_x,
            velocity_x + sound_speed,
        ],
        axis=0,
    )


def _eigen_lambdas_hydro(
    conserved_state,
    rhomin: Union[float, jnp.ndarray],
    pgmin: Union[float, jnp.ndarray],
    gamma: Union[float, jnp.ndarray],
    registered_variables: RegisteredVariables,
    mode: int,
):
    (
        velocity_x,
        sound_speed,
    ) = _eigenvalue_building_blocks(
        conserved_state,
        gamma,
        rhomin,
        pgmin,
        registered_variables,
    )

    def mode_0():
        return velocity_x - sound_speed

    def mode_1():
        return velocity_x

    def mode_2():
        return velocity_x

    def mode_3():
        return velocity_x

    def mode_4():
        return velocity_x + sound_speed

    return jax.lax.switch(
        mode, [mode_0, mode_1, mode_2, mode_3, mode_4]
    )