from dataclasses import dataclass, field
from enum import Enum
from typing import Literal, Union
import astropy.constants.astropyconst40 as c
import astropy.units as u
from math import pi
import numpy as np
from numpy import log, sqrt
from numpy.typing import NDArray
import scipy.constants as spc
from astropy import constants as const

from divergent_gravity.classes.bodys.laniakea import Laniakea


class DGDerivationStrategy(Enum):
    """
    If set to "G", the gravitational constant in standard SI units
    will be used to derive the cumulative universal mass.
    If set to "Mu", the standard value of 1.83x10^53 for the cumulative baryonic mass will be used to derive G.
    """

    G = "G"
    Mu = "Mu"


@dataclass
class DG:
    """
    Implementation of the Divergent Gravity (DG) model.
    This model treats gravity as a repulsive flux (c^2) and inertia
    as a Machian response to spatial density axis progression.
    """

    phi_infty: float = field(default_factory=lambda: spc.c**2)
    radius_hubble: float = field(
        default_factory=lambda: (14.4 * u.Glyr).decompose().value
    )
    hubble_param: float = 67.4
    A0: float = 0
    """
    A geometric scaling factor of A0.
    Defaults to 1/π
    """
    kappa: float = field(default_factory=lambda: 1 / pi)
    # H_infinity: float = field(default_factory=lambda: )
    Mu_observed: float = field(default_factory=lambda: 1.83 * 10**53)
    strategy: DGDerivationStrategy = field(
        default_factory=lambda: DGDerivationStrategy.G
    )
    laniakea: Laniakea = field(init=False)

    def __post_init__(self):
        # 4. Store as a float in SI units for your calculations
        self.A0 = self.get_machian_acceleration()

        # self.A0 = a0.to(u.m / u.s**2)
        if self.strategy == DGDerivationStrategy.G:
            self.laniakea = Laniakea(phi_infty=self.phi_infty, G=self.derived_G())
        else:
            self.laniakea = Laniakea(phi_infty=self.phi_infty)

    def phi_r(self, m: float = c.M_sun.value, r: Union[float, None] = None):
        G, _ = self.get_G_Mu()
        R = r if r is not None else spc.astronomical_unit
        return self.phi_infty - G * m / R

    def get_machian_acceleration(self):
        """
        Calculates the background acceleration threshold (a0).
        This is the gradient of the universal potential.
        """
        return self.phi_infty / self.radius_hubble

    def machian_coupling(self, M_local=2e47, R_local=2.3e24):
        """
        Final derivation of Omega_Mach (1.19) using Logarithmic Impedance Scaling.
        This represents the transition of the 4D density axis through
        the local 3D Laniakea shadow.
        """
        # 1. Calculate the dimensionless Impedance Ratio (Z_ratio)
        radius_ratio = R_local / self.radius_hubble
        _, Mu = self.get_G_Mu()
        mass_ratio = Mu / M_local

        # Z_ratio represents the volumetric displacement
        z_ratio = sqrt(mass_ratio * (radius_ratio**3))

        # 2. Logarithmic Mapping
        # Fields in a 4D divergent manifold scale logarithmically
        log_scale = log(z_ratio)

        # 3. Dimensional Projection Constant (pi / 2)
        # The ratio of the 4D rotation to the 3D surface progression
        projection_const = pi / 2

        return log_scale * projection_const

    def epsilon_0(self):
        """
        The electric constant. This right now is just a constant,
        but I'm leaving this method here for future iterations.
        """
        return spc.epsilon_0

    def mu_0(self):
        """
        The magnetic constant. This right now is just a constant,
        but I'm leaving this method here for future iterations.
        """
        return spc.mu_0

    def electron_radius(self, use_standard_electron_radius: bool = True):
        if use_standard_electron_radius:
            return spc.e**2 / (
                spc.electron_mass * spc.c**2 * 8 * pi * spc.epsilon_0
            )
        return (
            spc.e**2
            / (spc.electron_mass * self.phi_infty * 8 * pi * self.epsilon_0())
            * 2
        )

    def derived_G(self):
        return (self.phi_infty * self.radius_hubble) / self.Mu_observed

    def get_vacuum_impedance(self):
        """
        Returns the Impedance of Free Space (Z0).
        In DG, this represents the stiffness of the density axis.
        Standard value ~376.73 Ohms
        In a unified field, this is derived from the flux density.
        """
        return 120 * pi

    def derived_Mu(self):
        return self.phi_infty * (self.radius_hubble / spc.G)

    def get_G_Mu(self):
        G = spc.G if self.strategy == DGDerivationStrategy.G else self.derived_G()
        Mu = (
            self.Mu_observed
            if self.strategy == DGDerivationStrategy.Mu
            else self.derived_Mu()
        )
        return G, Mu

    def magnetic_divergence(self, volume_change_rate, kappa=1.0):
        """
        Calculates the predicted non-zero divergence of B field.
        grad.B = kappa * (dV/dt) / c^3
        """
        return (kappa * volume_change_rate) / (sqrt(self.phi_infty) ** 3)

    def g_tt(self, M_local, R_local):
        """
        Calculates the temporal metric component g_tt.
        In DG, this represents the local attenuation of the universal flux c^2.
        """
        G, _ = self.get_G_Mu()
        phi_local = (G * M_local) / R_local
        # Standard Schwarzschild-like component
        return 1 - (2 * phi_local / self.phi_infty)

    def alpha_inverse(self, M_local=2e47, R_local=2.3e24):
        """
        Full derivation of 1/alpha using 4D topological equilibrium.
        """
        # 1. Scale Gradient (Linear Density Progression)
        scale_gradient = log(self.radius_hubble / self.electron_radius())

        # 2. Hypersphere Gate Flux (Geometric Constant 3*pi^2)
        gate_flux = 3 * (pi**2)

        # 3. Angular Vorticity (Geometric Constant 4*pi)
        vorticity = 4 * pi

        # 4. Machian Coupling (The Environmental Metric Correction)
        mach_coupling = self.machian_coupling(M_local, R_local)

        alpha_inv = scale_gradient + gate_flux + vorticity + mach_coupling
        return alpha_inv

    # def predict_dilation_leak(self, v: Union[float, NDArray], acceleration_m_s2):
    #     """Completely theoretical."""
    #     a = acceleration_m_s2
    #     gamma = 1 / sqrt(1 - (v / sqrt(self.phi_infty)) ** 2)
    #     d_gamma_dt = (v * a * (gamma**3)) / (self.phi_infty)

    #     # Coupling constant of the 4D density axis
    #     kappa = (sqrt(self.mu_0() / self.epsilon_0())) / (4 * pi * self.radius_hubble)
    #     print("Kappa n dilation leak", kappa)
    #     return kappa * d_gamma_dt

    def summarize(self):
        """Prints the unified parameters of the model."""
        a0 = self.get_machian_acceleration()
        G, Mu = self.get_G_Mu()
        alpha_inv = self.alpha_inverse()

        print(
            f"""--- Divergent Gravity Model Summary ---
Hubble Radius (R_u):    {self.radius_hubble:2e} m
Universal Mass (M_u):   {Mu:.2e} kg
Machian Accel (a0):     {a0:.2e} m/s^2
Machian Coupling:       {self.machian_coupling():.6f}
Derived G:              {G:.2e} m^3/kg/s^2
Derived 1/alpha:        {alpha_inv:.4f}
---------------------------------------"""
        )
