import numpy as np
from numpy import sin, cos, tan, sqrt, exp, pi
import matplotlib.pyplot as plt
from scipy.constants import *  #This is love
from numba import jit
import time as time

solar_radius = 6.957*10**5
solar_mass = 1.9891*10**(30)
AU_to_km = 149597871
c_km = c/1000
rad_to_arcsec = 206264.80624709468

@jit(nopython = True)
def linearly_dt(x, y, t, N, index):  #Since I use a dynamic dt, it't necesarry to make it linearly when animating
    T = t[-1]
    line_time = np.linspace(0, T, N)
    x1, y1 = np.zeros(N), np.zeros(N)
    counter = 0
    flip = True
    for i in range(len(t)):
        if line_time[counter] <= t[i]:
            x1[counter] = x[i]
            y1[counter] = y[i]
            counter += 1
            if counter > N:
                break
            if line_time[counter] >= t[index]:
                if flip:
                    middle = counter
                    flip = False
    return x1, y1, middle

#Somewhat properly working animation
def animate(x, y, t, index, T = 45):  #Animating the light trajectory, #Work in progress
    """
    @T = The time for the total Animation,
    not in seconds but the lower the faster animation.
    """
    N = len(x)
    dt = 0.05
    total_frames = int(T/dt)

    x, y, middle = linearly_dt(x, y, t, total_frames, index)
    n = int(len(x)/4)

    x, y = np.copy(x[middle-n:middle+n]), np.copy(y[middle-n:middle+n])

    plt.ion()
    fig = plt.figure()
    ax = fig.add_subplot(111)
    line, = ax.plot( x[0], y[0], 'o', label='photon')
    ax.plot(x, y)

    fig.suptitle("Animation of gravitational lensing")
    ax.set_xlabel('Black hole radius')
    ax.set_ylabel('Black hole radius')
    ax.legend(loc='upper left')
    ax.set_xlim([x[0],x[-1]])
    ax.set_ylim([np.min(y), np.max(y)])
    plt.draw()

    for i in range(2*n):
        line.set_data(x[i], y[i])
        fig.canvas.draw()
        plt.pause(dt)

    plt.ioff()


@jit(nopython=True)
def pol2cart(r, phi):
    x = r*cos(phi)
    y = r*sin(phi)
    return x, y

@jit(nopython=True)
def dr(r, b):  #Light moving inwards
    r2 = r*r
    fact = (1-2/r)
    return fact*np.sqrt(1-fact*b*b/r2)

@jit(nopython=True)
def dphi(r, b):   #Light moving inwards
    r2 = r*r
    return b/r2*(1-2/r)

@jit(nopython=True)
def dr2(r, b):   #Light moving outwards
    fact = (1-2/r)
    return fact*np.sqrt(np.abs(1-fact*b**2/r**2))

@jit(nopython=True)
def dphi2(r, b):   #Light moving outwards
    return b/r**2*(1-2/r)

@jit(nopython=True)
def dt_select(r, x):  #Dynamic dt steps
    dt = r*x*r
    if dt > 100:
        return 100
    else:
        return dt

def wrapper(r, phi, index, index2, M):  #Simple wrapper that return both a aligned version of the lensing and the real trajectory
    try:
        x, y = np.copy(pol2cart(r[:index2], phi[:index2] - phi[index] + pi/2))
    except TypeError:
        x, y = np.copy(pol2cart(r, phi))

    try:
        x2, y2 = np.copy(pol2cart(r[:index2], phi[:index2]))
    except:
        x2, y2 = np.copy(pol2cart(r, phi))

    curvature = 4*M/(r[index])*rad_to_arcsec  #The analytical curvature
    num = 2*abs(abs(theta) - abs(np.arctan(y[-1]/x[-1])))*rad_to_arcsec  #The numerical curvature
    print('Analytical lensing angels: ', curvature)
    print('Numerical lensing angels: ', num)
    return x, y, x2, y2

@jit(nopython=True)
def grav_lens(r0, R_end, phi0, M, N=1e8):
    """
    @r0 = Distance from the black hole in AU
    R_end = The closest distance away from the black hole
    if you don't take into account gravity.
    Measured in solar radiuses.
    @phi0 = The start angle in polar coordinates of the light beam, measured in degrees.
    M = Mass of the black hole in solar masses
    N = lenght of the arrays, don't change
    """
    N = int(N)
    mass = M*solar_mass
    M = solar_mass*M*G/c**2/1000
    #T = 1.2*np.sqrt(r0**2 + R_end**2)*AU_to_km/c_km

    r0 = r0*AU_to_km/M
    R_end = R_end/M
    theta = np.arctan(R_end/r0)

    n = 4e7
    T = 2.4*np.sqrt(r0**2 + R_end**2)

    b = r0*np.sin(theta)
    r = np.zeros(N); r[0] = r0
    t = np.zeros(N)
    phi = np.zeros(N); phi[0] = phi0/180*pi
    flip = False
    counter = 0
    index2 = N-1

    #Not really a great solution for handeling diffrent distances but it works
    x = 0.002/R_end**2
    dt = dt_select(r0, x)
    x = dt/10*x

    for i in range(1, N):

        dt = dt_select(r[i-1], x)
        t[i] = t[i-1] + dt

        if flip == False:  ## TODO: Add the flipping inside the dr and dphi function
            drf1= dr(r[i-1], b)  #RK4
            dphif1 = dphi(r[i-1], b)  #The inwards movement equation
            x1 = r[i-1] - drf1*dt/2.0

            drf2 = dr(x1, b)
            dphif2 = dphi(x1, b)
            x2 = x1 - drf2*dt/2.0

            drf3 = dr(x2, b)
            dphif3 = dphi(x2, b)
            x3 = x2 - drf3*dt

            drf4 = dr(x3, b)
            dphif4 = dphi(x3, b)
            r[i] = r[i-1] - 1.0/6.0*(drf1+2*drf2+2*drf3+drf4)*dt
            phi[i] = phi[i-1] - 1.0/6.0*(dphif1+2*dphif2+2*dphif3+dphif4)*dt

            if np.isnan(r[i]):
                r[i] = r[i-1] - dr(r[i-1], b)*dt
                phi[i] = phi[i-1] - dphi(r[i], b)*dt

        else:
            drf1 = dr2(r[i-1], b)  #The outwards movement equation
            dphif1 = dphi2(r[i-1], b)
            x1 = r[i-1] + drf1*dt/2.0

            drf2 = dr2(x1, b)
            dphif2 = dphi2(x1, b)
            x2 = x1 + drf2*dt/2.0

            drf3 = dr2(x2, b)
            dphif3 = dphi2(x2, b)
            x3= x2 + drf3*dt

            drf4 = dr2(x3, b)
            dphif4 = dphi2(x3, b)
            r[i] = r[i-1] + 1.0/6.0*(drf1+2*drf2+2*drf3+drf4)*dt
            phi[i] = phi[i-1] - 1.0/6.0*(dphif1+2*dphif2+2*dphif3+dphif4)*dt

            if np.isnan(r[i]):
                r[i] = r[i-1] + dr2(r[i-1], b)*dt
                phi[i] = phi[i-1] - dphi2(r[i-1], b)*dt

            if r[i] > r0:  #Stopping the simulation when the light has traveled equal distance to and from the star
                index2 = i
                break

        if np.isnan(dr(r[i], b)):  #When the light trajectory turns, flipping which equations to use
            if flip == False:
                flip = True
                index = i

    if r[i] < r0:
        index2 = i
    print('Total iterations: ',index2)

    return r*M, phi, index, index2, theta, t[:index2], M  #Returning all in km

r, phi, index, index2, theta, t, M = grav_lens(4, 10*solar_radius, 180, 20000)  #Nice example of gravitational lensing
#r, phi, index, index2, theta, t, M = grav_lens(2, solar_radius, 180, 1)

x, y, x2, y2 = wrapper(r, phi, index, index2, M)

plt.title('Gravitational lensing of light, centered around the x-axis')
plt.plot([x[0], x[index]], [y[0], y[index]], color='green', label="Triangular trajectory")
plt.plot([x[index], x[0::100][-1]], [y[index], y[0::100][-1]], color='blue', label="Actual trajectory")
plt.plot([0,x[index]], [0, y[index]], linestyle = 'dashed', color='red')
plt.plot(x[index], y[index], 'ro', label='Closest point the the black hole')
plt.plot(x[0::100], y[0::100])
plt.legend()
plt.xlabel('[km]')
plt.ylabel('[km]')
plt.show()

plt.title('Gravitational lensing of light, actual trajectory')
plt.plot([x2[0], x2[index]], [y2[0], y2[index]], color='green', label="Triangular trajectory")
plt.plot([x2[index], x2[0::100][-1]], [y2[index], y2[0::100][-1]], color='blue', label = "Actual trajectory")
plt.plot([0,x2[index]], [0, y2[index]], linestyle = 'dashed', color='red')
plt.plot(x2[index], y2[index], 'ro', label='Closest point the the black hole')
plt.plot(x2[0::100], y2[0::100])
plt.legend()
plt.xlabel('[km]')
plt.ylabel('[km]')
plt.show()


animate(x2, y2, t, index)