'''This program calculates mass and density relationships for iron and carbon white dwarf stars
based upon a central density. This is done by using the balance of gravity
and pressure of relatavistic degenerate electron gas and solves using the fourth order
Runge-Kutta algorithm'''
#Norman Khan 2017

######################################
import math
from math import pi
import matplotlib.pyplot as plt
import numpy as np
#from uncertainties import ufloat

#####################################
##############CONSTANTS##############
#####################################
#Physical Constants
SPEED_OF_LIGHT=   2.99792458e8      #metres per second
H_BAR=            1.05457266e-34    #joule seconds
ELECTRON_CHARGE=  1.60217733e-19    #coulombs
ELECTRON_MASS=    9.1093897e-31     #kilograms
PROTON_MASS=      1.6726231e-27     #kilograms
GRAVITATION=      6.67259e-11       #MKS units
SOLAR_MASS=       1.98855e30        #kilograms
SOLAR_RADIUS=     6.957e8           #metres

#Known Stars
StarNames = ['Sirius B', '40 Eri B', 'Stein 2051']
StarMass = [1.053, 0.48, 0.50]          #Solar Masses
StarMassErr = [0.028, 0.02, 0.05]
StarRadius = [0.0074, 0.0124, 0.0115]   #Solar Radii
StarRadiusErr = [0.0006, 0.0005, 0.0012]
Markers = ['o', '*', '^']               #markers used on scatter plot

#####################################
##############VARIABLES##############
#####################################
CENTRAL_DENSITY=0.0                             #Density at the center of the white Dwarf

Y_e = 0.5                                       #Number of electrons per nucleon

Y_e_CARBON = 0.5                                #For Carbon Based White Dwarfs
Y_e_IRON = 26.0/56.0                            #For iron white Dwarfs


#For full explaination of these constants see associated paper.
RHO0 = 9.79e8 / Y_e                             #Bunch of constants over number of electrons per nucleon used for scaling density
RHO0_CARBON = 9.79e8 / Y_e_CARBON
RHO0_IRON = 9.79e8 / Y_e_IRON

R0 = ((3.0*Y_e*ELECTRON_MASS*SPEED_OF_LIGHT**2.0)/(4.0*pi*GRAVITATION*RHO0*PROTON_MASS))**0.5
R0_CARBON = ((3.0*Y_e_CARBON*ELECTRON_MASS*SPEED_OF_LIGHT**2.0)/(4.0*pi*GRAVITATION*RHO0_CARBON*PROTON_MASS))**0.5
R0_IRON = ((3.0*Y_e_IRON*ELECTRON_MASS*SPEED_OF_LIGHT**2.0)/(4.0*pi*GRAVITATION*RHO0_IRON*PROTON_MASS))**0.5

M0 = 4.0/3.0 * pi * RHO0 * R0**3                            #Constant used for scaling Mass (Total Mass)
M0_CARBON = 4.0/3.0 * pi * RHO0_CARBON * R0_CARBON**3       #Constant used for scaling Mass (Total Mass)
M0_IRON = 4.0/3.0 * pi * RHO0_IRON * R0_IRON**3             #Constant used for scaling Mass (Total Mass)

MASS, DENSITY, RADIUS = [], [], []                          #List for holding Mass, Density and Radius values
CENTRAL_DENSITY_LIST = []                                   #list for holding central densities

#Lists used for comparing stars/ groups of stars with different parameters
STAR_1_MASS, STAR_1_RADIUS, STAR_1_DENSITY = [], [], []
STAR_2_MASS, STAR_2_RADIUS, STAR_2_DENSITY = [], [], []

#####################################
##############FUNCTIONS##############
#####################################

###Gamma(y)### relativistic free Fermi gas
gamma = lambda y : (y**(2.0/3.0))/(3.0*(1.0+y**(2.0/3.0))**(1.0/2.0))


'''
x = independent variable   ----> r
y = list of dependent variables. eg [y1.y2]  --> [rho, m]
h = increment
n = order. eg 2 corresponds to two equations'''
def runkut(n, x, y, h):
    #"Advances the solution of diff eqn defined by derivs from x to x+h"
    y0=y[:]
    k1=derivs(n, x, y)
    for i in range(1,n+1): y[i]=y0[i]+0.5*h*k1[i]
    k2=derivs(n, x+0.5*h, y)
    for i in range(1,n+1): y[i]=y0[i]+h*(0.2071067811*k1[i]+0.2928932188*k2[i])
    k3=derivs(n, x+0.5*h, y)
    for i in range(1,n+1): y[i]=y0[i]-h*(0.7071067811*k2[i]-1.7071067811*k3[i])
    k4=derivs(n, x+h, y)
    for i in range(1,n+1):
        a=k1[i]+0.5857864376*k2[i]+3.4142135623*k3[i]+k4[i]
        y[i]=y0[i]+0.16666666667*h*a
    #print x, y
    x+=h
    return (x,y)


def derivs(n, x, y):
    #"The function DERIVS calculates y from x and y"
    dy=[0 for i in range(0,n+1)]
    #0 density is defined as the radius of the star
    if x==0:
        dy[1]=0
        return dy

    if y[1] >= 0:  #Only Calculate if density is > 0 (Ending boundry condition)
        dy[1] = -(y[2]*y[1])/((x**2)*gamma(y[1]))      #d'(rho)/dr'
        dy[2] = 3*(x**2)*y[1]                             #d'(m)/dr'
        return dy

    else: #if density is less than 0 then do not perform any calculations
        return dy


#####################################
##############   MAIN   #############
#####################################

#=========================================================#
#Mass vs Radius Iron and Carbon for many central densities#
#=========================================================#

#Calculating N_Stars Different Central Densities

#Plotting Many central densities Mass - Radius

N_Stars = 1000 #Number of stars of different central densities to create
for CENTRAL_DENSITY in np.logspace(-3, 6 ,N_Stars): #Full Radius Plot
#integrating from r=0 to r=R  with initial conditions rho(r) = CENTRAL_DENSITY   ,m(r) = 0
#r;   y=[ignore, rho, mass]

    N=2000 #number of interations per calculations for each star, higher number will result in higher resolution
    x=0.0; y=[0.0, CENTRAL_DENSITY, 0.0] # Set initial Boundary Conditions
    # Calculate and print the solution for x= 0 to 1 using N steps
    CENTRAL_DENSITY_INITIAL=CENTRAL_DENSITY
    #for j in range(0,N):
    while y[1]>1e-13:
        (x,y) = runkut(2, x, y, 1.0/N)

        MASS.append(y[2])
        RADIUS.append(x)
        DENSITY.append(y[1])

        #print "Central Density", CENTRAL_DENSITY,  "r =", x, "rho/rho0 =", y[1], "m/M0 =", y[2]
        if y[1]<1e-13:   #ending boundry condition rho(r)=0
            #print "rho =", y[1]
            #print "m =", y[2]
            STAR_1_MASS.append(MASS[-1]*M0_CARBON/SOLAR_MASS)
            STAR_1_RADIUS.append(RADIUS[-1]*R0_CARBON/SOLAR_RADIUS)
            STAR_1_DENSITY.append(CENTRAL_DENSITY_INITIAL)

            STAR_2_MASS.append(MASS[-1]*M0_IRON/SOLAR_MASS)
            STAR_2_RADIUS.append(RADIUS[-1]*R0_IRON/SOLAR_RADIUS)
            STAR_2_DENSITY.append(CENTRAL_DENSITY_INITIAL)

            #if 0.0115-0.0012 < STAR_2_RADIUS[-1] < 0.0115+0.0012:
            print  "Central density=", CENTRAL_DENSITY, "Radius=", STAR_2_RADIUS[-1], "mass=", STAR_2_MASS[-1]

            break

fig, ax1 = plt.subplots(figsize=(3.37,2.5))#10,8))
ax1.set_ylabel('Central Density ('+r'$\rho_{0}$)', fontsize=7)
ax1.set_xlabel('Mass ('+r'$M_{\odot}$)', fontsize=7)
ax1.tick_params(axis='both', which='major', labelsize=7)

'''
#Central density plot and axis
ax2 = ax1.twiny()
ax2.invert_xaxis()
ax2.tick_params(axis='both', which='major', labelsize=7)
#ax2.set_xscale('log')
ax2.set_xlabel('Central Density ('+r'$\rho_{0}$)', fontsize=7)
'''

'''
#Plotting known stars
for i in range(len(StarNames)):
    #ax1.annotate(StarNames[i], (StarRadius[i],StarMass[i]))
    #ax1.axvline(x=StarRadius[i], color='r', linewidth=0.25, linestyle='-')
    #ax1.text(StarRadius[i], 0.4, 'froggfhfghgfhgfhfhf', fontsize=4, rotation='vertical')

    ax1.scatter(StarRadius[i], StarMass[i], marker=Markers[i],s=10, lw = 0, color='k',  label=StarNames[i])
    ax1.errorbar(StarRadius[i], StarMass[i], xerr=StarRadiusErr[i], yerr=StarMassErr[i], capsize=1.0, elinewidth=0.25, color='k')
'''
ax1.plot(STAR_1_RADIUS, STAR_1_DENSITY, 'k', linewidth=0.5, label='Carbon White Dwarf')
ax1.plot(STAR_2_RADIUS, STAR_2_DENSITY,'k--', linewidth=0.5, label='Iron White Dwarf')

#ax2.plot(STAR_1_DENSITY, STAR_1_MASS, linewidth=1)
ax1.set_yscale('log')
ax1.set_ylim(ymin = 0.0)
ax1.set_xlim(xmin=0.000)
#plt.axvline(x=max(STAR_1_MASS), color='r', linewidth=0.5, linestyle='-')
#plt.axvline(x=max(STAR_2_MASS), color='r', linewidth=0.5, linestyle='-')

#ax1.text(max(STAR_1_MASS)-0.20, 0.20, r'%s $M_{\odot}$' %np.around(max(STAR_1_MASS), 2), fontsize=7)
#ax1.text(max(STAR_2_MASS)-0.25, 0.6, r'%s $M_{\odot}$' %np.around(max(STAR_2_MASS), 2), fontsize=7)

ax1.set_yscale('log')



ax1.legend(fontsize=7)
fig.savefig('Plot1.eps', format='eps', dpi=1000, bbox_inches='tight', pad_inches = 0.0)
fig.savefig('Plot1.png', format='png', dpi=1000, bbox_inches='tight', pad_inches = 0.0)


print "Carbon white dwarf maximum mass", STAR_1_MASS[-1]
print "Iron white dwarf maximum mass", STAR_2_MASS[-1]






######################################
#Plotting internal structure of stars#
######################################

#========================================================#
#Mass vs Radius Iron and Carbon for given central density#
#========================================================#

#resetting variables
CENTRAL_DENSITY = 1.0

fig2, ax2 = plt.subplots(figsize=(10,8))#figsize=(3.37,2.5))#
ax2.set_ylabel('Mass ('+r'$M_{\odot}$)', fontsize=7)
#ax2.set_ylabel('Density ('+r'$\rho_{0}$)', fontsize=7)
ax2.set_xlabel('Radius ('+r'$R_{\odot}$)', fontsize=7)
#plt.xticks(np.arange(0.0, 0.0016, 0.0005)) #Manual adjustment of xticks

ax2.tick_params(axis='both', which='major', labelsize=7) #Set x and y labels to fontsize 7



#while CENTRAL_DENSITY<1:

MASS, RADIUS, DENSITY = [],[],[]
STAR_1_MASS, STAR_1_RADIUS, STAR_1_DENSITY = [], [], []
STAR_2_MASS, STAR_2_RADIUS, STAR_2_DENSITY = [], [], []
N=2000 #number of interations per calculations for each star, higher number will result in higher resolution
x=0.0; y=[0.0, CENTRAL_DENSITY, 0.0] # Set initial Boundary Conditions


#for j in range(0,N):
while y[1]>1e-13:
    (x,y) = runkut(2, x, y, 1.0/N)# Calculate the solution for x= 0 to 1 using N steps
    #print x, y[1], y[2]
    MASS.append(y[2]*M0/SOLAR_MASS)
    RADIUS.append(x*R0/SOLAR_RADIUS)
    DENSITY.append(y[1])

    STAR_1_MASS.append(y[2]*M0_CARBON/SOLAR_MASS) #Carbon
    STAR_1_RADIUS.append(x*R0_CARBON/SOLAR_RADIUS)#Carbon
    STAR_1_DENSITY.append(y[1])

    STAR_2_MASS.append(y[2]*M0_IRON/SOLAR_MASS) #iron
    STAR_2_RADIUS.append(x*R0_IRON/SOLAR_RADIUS)#iron
    STAR_2_DENSITY.append(y[1])




ax2.plot(STAR_1_RADIUS,STAR_1_MASS,'k', linewidth=1.0, label='Carbon White Dwarf')
ax2.plot(STAR_2_RADIUS,STAR_2_MASS,'k--', linewidth=1.0, label='Iron White Dwarf')
print '------INDIVIDUAL STAR OF SINGULAR CENTRAL DENSITY----------'
print 'Central Density =',CENTRAL_DENSITY
print 'A carbon star has radius', STAR_1_RADIUS[-1], 'mass', STAR_1_MASS[-1]
print 'An Iron star has radius', STAR_2_RADIUS[-1], 'mass', STAR_2_MASS[-1]

'''
#Lines for Maximum Masses
plt.axhline(y=STAR_1_MASS[-1], color='r', linewidth=0.5, linestyle='-')
plt.axhline(y=STAR_2_MASS[-1], color='r', linewidth=0.5, linestyle='-')
ax2.text(0.00005, STAR_1_MASS[-1], r'%s $M_{\odot}$' %np.around(STAR_1_MASS[-1], 2), fontsize=7)
ax2.text(0.00005, STAR_2_MASS[-1], r'%s $M_{\odot}$' %np.around(STAR_2_MASS[-1], 2), fontsize=7)
'''

#Other annotations
ax2.text(0.0004, 0.35, r'Central Density: $\rho_{0}$', fontsize=7)
ax2.legend(fontsize=7)
plt.ylim(ymin=0)
plt.xlim(xmin=0)
#fig2.savefig('indv_Mass_Radius_rho0.eps', format='eps', dpi=1000, bbox_inches='tight', pad_inches = 0.0)