"""
Cosmological expansion dynamics via Friedmann equations.

This module implements self-consistent evolution of the scale factor a(τ)
in conformal time, coupled to field energy densities. The Friedmann equation
relates the Hubble parameter to the total energy density:

    H² = (8πG/3) ρ_total

In conformal time τ (related to physical time t by dt = a(τ)dτ), the
conformal Hubble parameter is H = a'/a.

For radiation-dominated universe: ρ ∝ a⁻⁴
For matter-dominated: ρ ∝ a⁻³

The solver uses a predictor-corrector scheme for numerical stability.
"""

import jax.numpy as jnp
from jax import jit
from typing import Tuple, Callable
from dataclasses import dataclass


@dataclass
class FRWUniverse:
    """
    FLRW universe parameters in conformal time.

    Attributes:
        a: Scale factor (dimensionless, typically normalized to a(τ₀) = 1)
        adot: Conformal time derivative da/dτ (same units as a/time)
        tau: Conformal time coordinate
        M_pl: Reduced Planck mass in simulation units (sets gravity strength)

    Derived quantities:
        H = adot/a: Conformal Hubble parameter
        rho_crit = 3H²/(8πG) = 3M_pl²H²: Critical density
    """

    a: float
    adot: float
    tau: float
    M_pl: float = 1.0  # Reduced Planck mass (default: natural units)

    @property
    def H(self) -> float:
        """Conformal Hubble parameter H = a'/a."""
        return self.adot / self.a if self.a > 0 else 0.0

    @property
    def rho_crit(self) -> float:
        """Critical density: 3H²/(8πG) = 3M_pl²H²/8π."""
        return 3.0 * self.M_pl**2 * self.H**2 / (8.0 * jnp.pi)

    def copy(self):
        """Create a copy of the universe state."""
        return FRWUniverse(a=self.a, adot=self.adot, tau=self.tau, M_pl=self.M_pl)

    def __repr__(self):
        return (
            f"FRWUniverse(a={self.a:.6f}, H={self.H:.6e}, "
            f"τ={self.tau:.4f}, M_pl={self.M_pl:.2f})"
        )


@jit
def friedmann_acceleration(rho: float, M_pl: float, a: float) -> jnp.ndarray:
    """
    Compute conformal acceleration a'' from Friedmann equation.

    From Friedmann: H² = (8πG/3)ρ with H = a'/a

    In conformal time:
        (a'/a)² = (8πG/3)ρ
        a'/a = √(8πG/3 * ρ)
        a' = a * √(8πG/3 * ρ)

    Taking derivative:
        a'' = a' * √(8πG/3 * ρ) + a * d/dτ[√(8πG/3 * ρ)]

    For simplicity (and typical radiation-dominated case where ρ ∝ a⁻⁴):
        a'' ≈ -H²a (deceleration from energy density)

    More accurately, using conservation ρ' = -4Hρ for radiation:
        a'' = (a'/a) * a' - a * (8πG/3) * ρ' / (2√(8πG/3 * ρ))

    Args:
        rho: Total energy density
        M_pl: Reduced Planck mass (G = 1/M_pl²)
        a: Current scale factor

    Returns:
        Conformal acceleration a''
    """
    # Friedmann equation: H = √(8πG/3 * ρ) = √(8π/(3M_pl²) * ρ)
    G = 1.0 / (M_pl * M_pl)
    H_squared = (8.0 * jnp.pi * G / 3.0) * rho
    H_squared_safe = jnp.maximum(H_squared, 0.0)  # Ensure non-negative for sqrt

    # For radiation-dominated: a'' = -H²a (standard deceleration)
    # This is exact for ρ ∝ a⁻⁴
    addot = -H_squared_safe * a

    return addot


def friedmann_step_predictor_corrector(
    universe: FRWUniverse,
    rho_func: Callable[[float], float],
    dt: float,
) -> FRWUniverse:
    """
    Advance scale factor by one timestep using predictor-corrector.

    This uses a 2nd-order accurate predictor-corrector scheme:
    1. Predictor: Euler step to estimate a(τ+dτ)
    2. Corrector: Trapezoidal rule using predicted value

    The energy density ρ is computed from field values via rho_func,
    which may depend on a (for rescaled fields).

    Args:
        universe: Current FRWUniverse state
        rho_func: Function computing total energy density from scale factor
                  Should have signature rho_func(a) -> rho
        dt: Conformal timestep dτ

    Returns:
        Updated FRWUniverse at τ + dτ
    """
    a = universe.a
    adot = universe.adot
    tau = universe.tau
    M_pl = universe.M_pl

    # Current energy density and acceleration
    rho_current = rho_func(a)
    addot_current = friedmann_acceleration(rho_current, M_pl, a)

    # Predictor step (Euler) with positivity clamp
    a_pred_val = float(a + dt * adot)
    a_pred_val = max(a_pred_val, 1e-12)
    adot_pred = adot + dt * addot_current

    # Evaluate at predicted point
    rho_pred = rho_func(a_pred_val)
    addot_pred = friedmann_acceleration(rho_pred, M_pl, a_pred_val)

    # Corrector step (trapezoidal rule)
    a_new_val = float(a + dt * 0.5 * (adot + adot_pred))
    a_new_val = max(a_new_val, 1e-12)
    adot_new = adot + dt * 0.5 * (addot_current + addot_pred)
    tau_new = tau + dt

    return FRWUniverse(a=a_new_val, adot=float(adot_new), tau=tau_new, M_pl=M_pl)


def friedmann_step_leapfrog(
    universe: FRWUniverse,
    rho_func: Callable[[float], float],
    dt: float,
) -> FRWUniverse:
    """
    Advance scale factor using leapfrog integration (symplectic).

    Leapfrog for a(τ):
        a'(τ+dτ/2) = a'(τ) + (dτ/2) a''(τ)
        a(τ+dτ) = a(τ) + dτ a'(τ+dτ/2)
        a'(τ+dτ) = a'(τ+dτ/2) + (dτ/2) a''(τ+dτ)

    This preserves symplectic structure if rho_func is Hamiltonian.

    Args:
        universe: Current FRWUniverse state
        rho_func: Function computing ρ(a)
        dt: Conformal timestep

    Returns:
        Updated FRWUniverse
    """
    a = universe.a
    adot = universe.adot
    tau = universe.tau
    M_pl = universe.M_pl

    # Half-step velocity
    rho = rho_func(a)
    addot = friedmann_acceleration(rho, M_pl, a)
    adot_half = adot + 0.5 * dt * addot

    # Full step position
    a_new = a + dt * adot_half

    # Half-step velocity (second half)
    rho_new = rho_func(a_new)
    addot_new = friedmann_acceleration(rho_new, M_pl, a_new)
    adot_new = adot_half + 0.5 * dt * addot_new

    tau_new = tau + dt

    return FRWUniverse(a=a_new, adot=adot_new, tau=tau_new, M_pl=M_pl)


def radiation_scaling_check(
    a_initial: float,
    a_final: float,
    rho_initial: float,
    rho_final: float,
) -> Tuple[float, float]:
    """
    Check if energy density scales as radiation (ρ ∝ a⁻⁴).

    For pure radiation:
        ρ(a) = ρ₀ (a₀/a)⁴

    Returns:
        scaling_exponent: Measured exponent n in ρ ∝ a⁻ⁿ
        relative_error: Deviation from n=4
    """
    if a_initial <= 0 or a_final <= 0 or rho_initial <= 0 or rho_final <= 0:
        return float("nan"), float("nan")

    # From ρ_f/ρ_i = (a_i/a_f)^n, solve for n:
    # log(ρ_f/ρ_i) = n log(a_i/a_f)
    # n = log(ρ_f/ρ_i) / log(a_i/a_f)

    ratio_rho = rho_final / rho_initial
    ratio_a = a_initial / a_final

    if ratio_a <= 0:
        return float("nan"), float("nan")

    n = jnp.log(ratio_rho) / jnp.log(ratio_a)

    expected_n = 4.0  # Radiation
    relative_error = abs(n - expected_n) / expected_n

    return float(n), float(relative_error)


def matter_scaling_check(
    a_initial: float,
    a_final: float,
    rho_initial: float,
    rho_final: float,
) -> Tuple[float, float]:
    """
    Check if energy density scales as matter (ρ ∝ a⁻³).

    Returns:
        scaling_exponent: Measured n in ρ ∝ a⁻ⁿ
        relative_error: Deviation from n=3
    """
    if a_initial <= 0 or a_final <= 0 or rho_initial <= 0 or rho_final <= 0:
        return float("nan"), float("nan")

    ratio_rho = rho_final / rho_initial
    ratio_a = a_initial / a_final

    if ratio_a <= 0:
        return float("nan"), float("nan")

    n = jnp.log(ratio_rho) / jnp.log(ratio_a)

    expected_n = 3.0  # Matter
    relative_error = abs(n - expected_n) / expected_n

    return float(n), float(relative_error)


def create_radiation_universe(
    a_initial: float = 1.0,
    rho_initial: float = 1.0,
    M_pl: float = 1.0,
) -> FRWUniverse:
    """
    Initialize FRW universe in radiation-dominated era.

    For radiation: H² = (8πG/3)ρ with ρ ∝ a⁻⁴
    Initial velocity: a' = a H

    Args:
        a_initial: Initial scale factor
        rho_initial: Initial energy density
        M_pl: Reduced Planck mass

    Returns:
        Initialized FRWUniverse
    """
    G = 1.0 / (M_pl * M_pl)
    H_squared = (8.0 * jnp.pi * G / 3.0) * rho_initial
    H = jnp.sqrt(H_squared)

    adot_initial = float(a_initial * H)

    return FRWUniverse(a=a_initial, adot=adot_initial, tau=0.0, M_pl=M_pl)


def create_matter_universe(
    a_initial: float = 1.0,
    rho_initial: float = 1.0,
    M_pl: float = 1.0,
) -> FRWUniverse:
    """
    Initialize FRW universe in matter-dominated era.

    For matter: H² = (8πG/3)ρ with ρ ∝ a⁻³

    Args:
        a_initial: Initial scale factor
        rho_initial: Initial energy density
        M_pl: Reduced Planck mass

    Returns:
        Initialized FRWUniverse
    """
    G = 1.0 / (M_pl * M_pl)
    H_squared = (8.0 * jnp.pi * G / 3.0) * rho_initial
    H = jnp.sqrt(H_squared)

    adot_initial = float(a_initial * H)

    return FRWUniverse(a=a_initial, adot=adot_initial, tau=0.0, M_pl=M_pl)
