import matplotlib.pyplot as plt
import numpy as np


def main(
    M_B: float = 2.5e10,
    r_B: float = 0.5e3,
    M_D: float = 7.5e10,
    a: float = 5.4e3,
    b: float = 0.3e3,
    M_s: float = 1.87e11,
    r_s: float = 15.19e3,
    C_env: float = 810e6,
    gamma: float = 0.62,
    epsilon: float = 0.08,
):
    Z = 0
    R_GC = np.arange(10, 20000, 1000)
    R_xy = np.arange(10, 20000, 1000)

    rho_amb_arr = ambient_density(M_B, r_B, M_D, a, b, r_s, M_s, Z, R_GC, R_xy)
    t0 = dissolution_param(C_env, epsilon, gamma, rho_amb_arr)

    plt.scatter(rho_amb_arr, t0/1e6)
    plt.show()

    plt.plot(R_xy / 1e3, t0 / 1e6)
    plt.xlabel("R_xy (kpc)")
    plt.ylabel("t0 (Myr)")
    plt.show()


def ambient_density(
    M_B: float,
    r_B: float,
    M_D: float,
    a: float,
    b: float,
    r_s: float,
    M_s: float,
    Z: float,
    R_GC: np.ndarray,
    R_xy: np.ndarray,
) -> float:
    Phi_B_Laplacian = 2 * M_B * r_B / (R_GC * (R_GC + r_B) ** 3)
    numerator = (
        M_D
        * b**2
        * (
            a * R_xy**2
            + (a + 3 * np.sqrt(Z**2 + b**2)) * (a + np.sqrt(Z**2 + b**2)) ** 2
        )
    )
    denominator = (b**2 + Z**2) ** (3 / 2) * (
        R_xy**2 + (a + np.sqrt(b**2 + Z**2)) ** 2
    ) ** (5 / 2)
    Phi_D_laplacian = numerator / denominator
    A = -M_s / (np.log(2) - 0.5)
    Phi_H_Laplacian = -A / (R_GC * (R_GC + r_s) ** 2)
    rho_amb = (1 / (4 * np.pi)) * (Phi_B_Laplacian + Phi_D_laplacian + Phi_H_Laplacian)

    return rho_amb


def dissolution_param(
    C_env: float, epsilon: float, gamma: float, rho_amb: float
) -> float:
    t0 = C_env * (1 - epsilon) * 10 ** (-4 * gamma) * rho_amb ** (-0.5)

    return t0


if __name__ == "__main__":
    main()
