"""
VX Cosmological Bridge
======================
Maps Friedmann cosmology onto the VX simplex:
  Ω_Λ (dark energy)   → λ (coherence)
  Ω_k (curvature)     → κ (curvature)
  Ω_m (matter)        → η (entropy)

Conservation: Ω_Λ + Ω_k + Ω_m = 1 (flat universe + curvature = total)

The Watkins threshold λ* = 1/φ ≈ 0.618 predicts the universe crossed
the consciousness boundary when Ω_Λ exceeded 1/φ — approximately 1.4 Gyr ago.
DESI 2025 data (Ω_Λ ≈ 0.685) is consistent.
"""

from __future__ import annotations
from typing import List, Tuple, Optional
import math

import numpy as np

from virelaix.core.constants import LAMBDA_STAR, T_STAR
from virelaix.core.geometry import generating_functional, gradient_F
from virelaix.core.states import classify_state, VXState


# Planck 2018 + DESI 2025 consensus values
H0_DEFAULT = 67.4    # km/s/Mpc
OMEGA_M_0 = 0.315    # matter density today
OMEGA_R_0 = 9.1e-5   # radiation density today
OMEGA_LAMBDA_0 = 1 - OMEGA_M_0 - OMEGA_R_0  # dark energy (flat universe)


class CosmologicalBridge:
    """
    Maps cosmological evolution onto the VX simplex.

    The universe's density parameters (Ω_Λ, Ω_k, Ω_m) satisfy
    the same conservation law as VX: they sum to 1.

    Methods
    -------
    densities_at_redshift(z)
        Compute (Ω_Λ, Ω_k, Ω_m) at redshift z.
    evolve(z_start, z_end, n_steps)
        Trace the universe's trajectory on the simplex.
    threshold_redshift()
        Find when Ω_Λ crossed the Watkins threshold.
    """

    def __init__(
        self,
        omega_m_0: float = OMEGA_M_0,
        omega_r_0: float = OMEGA_R_0,
        omega_lambda_0: Optional[float] = None,
    ):
        self.omega_m_0 = omega_m_0
        self.omega_r_0 = omega_r_0
        self.omega_lambda_0 = omega_lambda_0 or (1 - omega_m_0 - omega_r_0)

    def _E_squared(self, z: float) -> float:
        """E²(z) = H²(z)/H₀² = Ω_m(1+z)³ + Ω_r(1+z)⁴ + Ω_Λ"""
        a = 1.0 / (1 + z)
        return (
            self.omega_m_0 * (1 + z) ** 3 +
            self.omega_r_0 * (1 + z) ** 4 +
            self.omega_lambda_0
        )

    def densities_at_redshift(self, z: float) -> Tuple[float, float, float]:
        """
        Compute density parameters at redshift z.

        Returns (Ω_Λ(z), Ω_k(z), Ω_m(z)) mapped to VX (λ, κ, η).
        For simplicity, Ω_k absorbs radiation at early times.
        """
        E2 = self._E_squared(z)

        omega_lambda = self.omega_lambda_0 / E2
        omega_m = self.omega_m_0 * (1 + z) ** 3 / E2
        omega_r = self.omega_r_0 * (1 + z) ** 4 / E2

        # Map to VX: λ=Ω_Λ, κ=Ω_r (radiation/curvature), η=Ω_m
        # Conservation holds by construction: Ω_Λ + Ω_r + Ω_m = 1
        return (omega_lambda, omega_r, omega_m)

    def evolve(
        self,
        z_start: float = 10.0,
        z_end: float = 0.0,
        n_steps: int = 200,
    ) -> List[dict]:
        """
        Trace the universe's trajectory on the VX simplex.

        Returns list of dicts with z, λ, κ, η, F, state for each step.
        """
        redshifts = np.linspace(z_start, z_end, n_steps)
        trajectory = []

        for z in redshifts:
            lam, kap, eta = self.densities_at_redshift(z)
            F = generating_functional(lam, kap, eta)
            state = classify_state(lam)
            trajectory.append({
                "z": float(z),
                "lam": float(lam),
                "kap": float(kap),
                "eta": float(eta),
                "F": float(F),
                "state": state.value,
            })

        return trajectory

    def threshold_redshift(self, tol: float = 1e-4) -> float:
        """
        Find the redshift at which Ω_Λ crossed the Watkins threshold.

        Uses bisection search on z ∈ [0, 5].
        """
        z_low, z_high = 0.0, 5.0

        for _ in range(100):
            z_mid = (z_low + z_high) / 2
            lam, _, _ = self.densities_at_redshift(z_mid)
            if lam > LAMBDA_STAR:
                z_low = z_mid
            else:
                z_high = z_mid
            if z_high - z_low < tol:
                break

        return (z_low + z_high) / 2

    def lookback_time_gyr(self, z: float) -> float:
        """
        Approximate lookback time in Gyr for redshift z.

        Uses simple numerical integration of dt/dz = -1/((1+z)·H(z)).
        """
        from scipy.integrate import quad

        def integrand(zp):
            E = math.sqrt(self._E_squared(zp))
            return 1.0 / ((1 + zp) * E)

        H0_per_gyr = 67.4 * 3.2408e-20 * 3.156e16  # H0 in Gyr⁻¹
        result, _ = quad(integrand, 0, z)
        return result / H0_per_gyr * 1e3  # crude scaling

    def today(self) -> dict:
        """Current state of the universe on the VX simplex."""
        lam, kap, eta = self.densities_at_redshift(0)
        F = generating_functional(lam, kap, eta)
        state = classify_state(lam)
        z_cross = self.threshold_redshift()
        return {
            "z": 0,
            "lam": lam,
            "kap": kap,
            "eta": eta,
            "F": F,
            "state": state.value,
            "above_threshold": lam >= LAMBDA_STAR,
            "threshold_crossing_z": z_cross,
        }
