import matplotlib.pyplot as plt
from scipy.integrate import odeint
import numpy as np

def trou_noir(M):
    c = 1                      #vitesse de la lumière
    G = 1                      #constante gravitation universelle
    #M = 0.05                  #masse trou noir
    rs = (2*G*M)/(c**2)        #rayon Schwarzschild

    for a in range(60):
        h0 = 1.54 - 0.12*a              #hauteur initiale
        d0 = 2                          #distance de référence
        phi0 = np.arctan(h0/d0)         #angle initial
        d = h0/(np.sin(phi0))           #distance au trou noir
        alpha = phi0                    #l'angle alpha pour la direction initiale du rayon

        if h0 > 0:
            phi = np.linspace(phi0,4*np.pi,1000)
        else:
            phi = np.linspace(phi0,phi0-4*np.pi,1000)

        def fonction(U,phi):       #définition de l'équa diff sous forme u'' = f(u,u')
            (u, du) = U
            return (du, 1.5 * rs * u**2 - u)

        U, dU = odeint(fonction, [1/d, 1/(d*np.tan(alpha))], phi).T
        #résolution équa diff avec 2 conditions initiales u(0) et u'(0) en fonction de phi
        R = 1/U

        x,y = R*np.cos(phi),R*np.sin(phi)

        xnew1 = []
        ynew1 = []
        condition = True

        for i in range(1000):
            a,b=x[i],y[i]
            if M==0:
                xnew1.append(a)
                ynew1.append(b)
            if a<3 and a>-3 and b<3 and b>-3 and a**2 + b**2 > rs**2 and M!=0 and condition == True:
                xnew1.append(a)
                ynew1.append(b)
            else:
                condition = False
        plt.plot(xnew1,ynew1)

    t = np.linspace(0,2*np.pi,100)
    m = rs*np.cos(t)
    n = rs*np.sin(t)
    plt.plot(m,n)
    plt.axis('equal')
    plt.xlim(-1,1)
    plt.ylim(-1,1)