using LinearAlgebra

function kerr_newman_metric(M, a, Q, r, θ)
    ρ² = r^2 + a^2*cos(θ)^2     # The horizon and singularity function
    Δ = r^2 - 2M*r + a^2 + Q^2  # Discriminant for event horizons
    
    g = zeros(4,4)
    g[1,1] = -(Δ - a^2*sin(θ)^2)/ρ²                      # g_tt - temporary component
    g[1,4] = g[4,1] = -a*sin(θ)^2*(2M*r - Q^2)/ρ²        # g_tϕ - rotation component
    g[2,2] = ρ²/Δ                                        # g_rr - radial component
    g[3,3] = ρ²                                          # g_θθ - angular component
    g[4,4] = ((r^2+a^2)^2 - Δ*a^2*sin(θ)^2)*sin(θ)^2/ρ²  # g_ϕϕ - azimuthal component
    
    return g
end

function compute_everything(M, a, Q, r, θ)
    # All in one function
    g = kerr_newman_metric(M, a, Q, r, θ)  # The metric tensor
    
    # Four-vector speed for a static observer
    u = zeros(4); u[4] = 1.0  # Initial condition: dt/dt = 1
    
    # Normalization: g_μν u^μ u^v = -1 (for a time-like vector)
    norm_squared = dot(u, g, u)  # This should be negative for a time-like vector.
    if norm_squared >= 0
        # If the vector is spatially similar, this is a problem - we use safe normalization.
        u ./= sqrt(abs(norm_squared))
    else
        u ./= sqrt(-norm_squared)  # The root of a positive number
    end
    
    # A simple calculation of the curvature through the second derivative
    h = 1e-6  # A step for numerical differentiation
    g_rp = kerr_newman_metric(M, a, Q, r+h, θ)  # The metric at the point r+h
    g_rm = kerr_newman_metric(M, a, Q, r-h, θ)  # The metric at the r-h point
    R_simple = (g_rp[1,1] - 2g[1,1] + g_rm[1,1]) / h^2  # The final difference of the second derivative of g_tt
    
    return g, u, R_simple
end

# Parameters of the black hole and coordinates of the observation point
M = 1.0    # The mass of the black hole in natural units
a = 0.5    # Rotation parameter (a = J/M)
Q = 0.3    # The electric charge of a black hole
r = 2.0    # Radial coordinate (distance from the center)
θ = π/4    # Angular coordinate (π/4 = 45°)

# Calculation of all values at a given point
g, u, R = compute_everything(M, a, Q, r, θ)

# Output of results
println("Kerr-Newman metric at the point (r=$r, θ=$θ):")
display(g)
println("\The four speed vector of the static observer:")
println("u = $u")  # Components: [u^r, u^θ, u^ϕ, u^t]
println("Simplified measure of curvature (second derivative of g_tt):")
println("R ≈ $R")

using Plots.PlotMeasures

# Creating a grid for visualization
r_range = range(1.5, 5.0, length=50)  # Radius from 1.5 to 5 (avoiding the singularity)
θ_range = range(0.1, π-0.1, length=50)  # Angle from 0.1 to π-0.1 (avoiding the pole)

# Calculating the metric and curvature components on the grid
g_tt_values = [kerr_newman_metric(M, a, Q, r, θ)[1,1] for r in r_range, θ in θ_range]
g_rr_values = [kerr_newman_metric(M, a, Q, r, θ)[2,2] for r in r_range, θ in θ_range]
curvature_values = [compute_everything(M, a, Q, r, θ)[3] for r in r_range, θ in θ_range]

p1 = heatmap(r_range, θ_range, g_tt_values', 
            title="The time component of the g_tt metric",
            ylabel="Angle θ",
            color=:viridis, clim=(-2, 0), titlefont=font(12), topmargin=2mm, titlelocation=:left)

p2 = heatmap(r_range, θ_range, g_rr_values', 
            title="The radial component of the g_rr metric", 
            ylabel="Angle θ",
            color=:plasma, clim=(0, 10), titlefont=font(12), topmargin=2mm, titlelocation=:left)

p3 = heatmap(r_range, θ_range, curvature_values',
            title="A simplified measure of curvature",
            xlabel="Radius r", ylabel="Angle θ", 
            color=:hot, clim=(-1, 1), titlefont=font(12), topmargin=2mm, titlelocation=:left)

# Displaying all the graphs together
plot(p1, p2, p3, layout=(3,1), size=(800, 900))