import numpy as np
import plotly.graph_objects as go
import plotly.io as pio

pio.templates.default = "plotly_dark"


class SPARCValidator:
    def __init__(self, sparc_data):
        self.data = sparc_data
        self.A0 = 8.740391e-10  # Pioneer-consistent Machian Anchor
        self.KAPPA_BASE = 1.0 / np.pi
        self.G = 6.67430e-11
        self.KPC_TO_M = 3.086e19

    def estimate_half_mass_radius(self, r_kpc, v_gas, v_disk, v_bulge):
        """
        Step 1: Calculate the baryonic mass profile.
        Step 2: Find the radius containing 50% of total mass.
        """
        # Sum of baryonic components (Newtonian sum)
        v_bar_sq = np.abs(v_gas) * v_gas + v_disk**2 + v_bulge**2

        # M(r) in kg = (V^2 * r) / G
        m_r = (v_bar_sq * (r_kpc * self.KPC_TO_M)) / self.G

        # Guard against noisy data at the core
        m_r = np.maximum.accumulate(m_r)

        m_total = m_r[len(m_r) - 1]
        if m_total <= 0:
            return 3.5  # Fallback to average spiral scale

        # Find index where mass crosses the 50% threshold
        half_mass_idx = np.argmin(np.abs(m_r - (m_total * 0.5)))
        return r_kpc[half_mass_idx]

    def calculate_gdg(self, r_kpc, v_gas, v_disk, v_bulge):
        # 1. Total Newtonian Baryonic Velocity
        v_bar_sq = np.abs(v_gas) * v_gas + v_disk**2 + v_bulge**2
        v_bar = np.sqrt(np.maximum(v_bar_sq, 0))

        # 2. Estimate the Machian Aperture (R_1/2)
        r_half_kpc = self.estimate_half_mass_radius(r_kpc, v_gas, v_disk, v_bulge)

        # 3. Conversions to SI
        r_m = r_kpc * self.KPC_TO_M
        r_half_m = r_half_kpc * self.KPC_TO_M
        gn = (v_bar * 1000.0) ** 2 / r_m
        # Avoid division by zero at r=0
        gn = np.where(r_m == 0, 1e-20, gn)

        # 4. Scaled Kappa Correction: (r / (r + r_half))
        # This modulates the Machian flux interaction
        kappa_r = self.KAPPA_BASE * (r_m / (r_m + r_half_m))

        # 5. DG Identity: g_dg = sqrt(gn^2 + gn * A0 * kappa)
        gdg = np.sqrt(gn**2 + gn * self.A0 * kappa_r)

        return np.sqrt(gdg * r_m) / 1000.0

    def validate_galaxy(self, g_id):
        g = self.data[g_id]
        v_pred = self.calculate_gdg(g["Rad"], g["Vgas"], g["Vdisk"], g["Vbul"])

        # Calculate residuals
        # We handle cases where v_obs might be zero or NaN
        mask = (g["Vobs"] > 0) & (~np.isnan(v_pred))
        mape = np.mean(np.abs((g["Vobs"][mask] - v_pred[mask]) / g["Vobs"][mask])) * 100

        return v_pred, mape


def generate_residual_map(validator, sparc_database):
    galaxy_ids = []
    mapes = []
    surface_brightness = []
    half_mass_radii = []

    # 1. Process the database
    for g_id in sparc_database.keys():
        v_pred, mape = validator.validate_galaxy(g_id)

        g = sparc_database[g_id]
        if len(g["Vdisk"]) > 0 and np.max(g["Vdisk"]) > 0:
            v_max = np.max(g["Vdisk"])
            r_peak = g["Rad"][np.argmax(g["Vdisk"])]
            sigma_proxy = (v_max**2) / (r_peak)

            # Use our validator's internal method to get the R_1/2 used
            r_half = validator.estimate_half_mass_radius(
                g["Rad"], g["Vgas"], g["Vdisk"], g["Vbul"]
            )

            galaxy_ids.append(g_id)
            mapes.append(round(mape, 6))
            surface_brightness.append(sigma_proxy)
            half_mass_radii.append(round(r_half, 2))

    print(f"Mean MAPE: {np.array(mapes).mean()}")

    # Plotting the Final Verdict
    # 2. Create the Plotly Figure
    fig = go.Figure()

    # Add the scatter points
    fig.add_trace(
        go.Scatter(
            x=surface_brightness,
            y=mapes,
            mode="markers",
            marker=dict(
                size=10,
                color=mapes,  # Color by error magnitude
                colorscale="Viridis",
                showscale=True,
                colorbar=dict(title="MAPE (%)"),
                line=dict(width=1, color="DarkSlateGrey"),
            ),
            text=galaxy_ids,  # Base text for hover
            # Customizing the hover tooltip
            hovertemplate=(
                "<b>Galaxy: %{customdata[0]}</b><br>"
                + "MAPE: %{y}%<br>"
                + "Sigma Proxy: %{x:.2f}<br>"
                + "Half-Mass Radius: %{customdata[1]} kpc"
                + "<extra></extra>"  # Removes the 'trace 0' label
            ),
            customdata=np.stack((galaxy_ids, half_mass_radii), axis=-1),
        )
    )

    # 3. Add the "Paper Boundary" Horizontal Line
    fig.add_shape(
        type="line",
        x0=min(surface_brightness),
        x1=max(surface_brightness),
        y0=1.0,
        y1=1.0,
        line=dict(color="Red", width=2, dash="dash"),
    )

    # 4. Styling the Layout
    fig.update_layout(
        title="Divergent Gravity Scale Invariance: Global SPARC Residuals",
        xaxis_title="Central Surface Brightness Proxy (Log Scale)",
        yaxis_title="Model MAPE (%)",
        xaxis_type="log",
        # template="plotly_white",
        hoverlabel=dict(bgcolor="white", font_size=12),
    )

    fig.show()

    return np.mean(mapes)
