from dataclasses import dataclass, field
from enum import Enum
from typing import 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 dg.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):
        # 1. Properly define H0 as a frequency (1/s)
        h0_astronomy = self.hubble_param * (u.km / u.s / u.Mpc)
        h0_frequency = h0_astronomy.to(u.s**-1)

        # 2. Ensure phi_infty has velocity-squared units (m^2/s^2)
        # Taking the sqrt of phi_infty gives us 'c' in m/s
        c_speed = np.sqrt(self.phi_infty) * (u.m / u.s)

        # 3. Calculate acceleration: a0 = c * H0
        # Units: (m/s) * (1/s) = m/s^2
        a0_quantity = c_speed * h0_frequency

        # 4. Store as a float in SI units for your calculations
        self.A0 = a0_quantity.to(u.m / u.s**2).value * self.kappa

        # 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 calculate_flat_velocity(self, M_galaxy):
        """
        Directly derives flat orbital velocity using the
        Mass Ratio Identity (v^4 = (M/Mu) * c^4).
        """
        _, Mu = self.get_G_Mu()
        # This is the 'Pure' DG derivation that bypasses G entirely.
        v_fourth = (M_galaxy / Mu) * (self.phi_infty**2)
        return pow(v_fourth, 0.25)

    def predict_dilation_leak(self, v: Union[float, NDArray], acceleration_m_s2):
        """Predicts the magnetic divergence caused by spatial dilation."""
        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)
        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}
        ---------------------------------------"""
        )
