# Play with these parameters. Lengths normalised to R0=1
# Epsilon is inverse aspect ratio, kappa is elongation, delta is triangularity, and d(p)/d(psi)= -C/4pi, d(F^2)/d(psi) = 2A. 
# C=1+A fixes the normalisation of psi0.
# external_coil_current does not affect psi AT ALL, but affects the toroidal beta. Needs to be larger than poloidal currents in plasma, otherwise toroidal field will vanish somewhere!

# TFTR
epsilon = 1/2.9
kappa = 1
delta = 0
A = 1
external_coil_current = 1.5 

# NSTX
#epsilon = 0.78
#kappa = 2
#delta = 0.35
#A = 0.05
#external_coil_current = 2.3

# DIII-D
#epsilon = 0.37
#kappa = 2.1
#delta = -0.5
#A = 0.5
#external_coil_current = 1.5

# ITER
#epsilon = 0.32
#kappa = 1.7
#delta = 0.33
#A = 0.155
#external_coil_current = 1.25

# JET
#epsilon = 1/3
#kappa = 1.7
#delta = 0.25
#A = 0.2
#external_coil_current = 1.15

# MAST
#epsilon = 1/1.3
#kappa = 2.45
#delta = 0.5
#A = 0.05
#external_coil_current = 1.1


C = 1+A    

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import sympy as sp
sig_figs = 4

print('-------------------------------------------------')
print(f"Grad-Shafranov Solver using Solov'ev Profiles")
print('-------------------------------------------------\n')
print(f"Lengths normalised to 'average' major radius R0")
print(f"Parameters are")
print(f"\tinverse aspect ratio = {epsilon:.{sig_figs}g}")
print(f"\telongation = {kappa:.{sig_figs}g}")
print(f"\ttriangularity = {delta:.{sig_figs}g}")
print(f"\tpressure profile is 4*pi*p'(psi) = -{C}")
print(f"\tpoloidal current profile is F(psi)^2 = {external_coil_current**2:.{sig_figs}g} + {2*A}psi\n")


def solve_solovev_grad_shafranov(epsilon,kappa,delta_thing,alpha,psi0):
    #Following 6.6 of Freidberg, this code defines the particular solution and basis functions for the complementary solution of the G-S equation.
    delta = np.arcsin(delta_thing)
    r = sp.Symbol('r')
    z = sp.Symbol('z')
    up = -1*alpha/2 * r**2 * sp.log(r) + (1+alpha)/8 * r**4
    us= [
        r**0,
        r**2,
        z**2 - r**2 * sp.log(r),
        r**4-4* r**2 * z**2,
        2 * z**4 - 9 * z**2*r**2 - (12*z**2* r**2 - 3*r**4)*sp.log(r),
        r**6 - 12*r**4*z**2 + 8*r**2*z**4,
        8*z**6-140*z**4*r**2+75*z**2*r**4-(120*z**4*r**2-180*z**2*r**4+15*r**6)*sp.log(r)
    ]
    up_r = sp.diff(up,r)
    up_rr = sp.diff(up_r,r)
    up_z = sp.diff(up,z)
    up_zz = sp.diff(up_z,z)
    us_r = [sp.diff(u,r) for u in us]
    us_rr = [sp.diff(u_r,r) for u_r in us_r]
    us_z = [sp.diff(u,z) for u in us]
    us_zz = [sp.diff(u_z,z) for u_z in us_z]
    func_up = sp.lambdify((r,z),up,'numpy')
    func_up_r = sp.lambdify((r,z),up_r,'numpy')
    func_up_rr = sp.lambdify((r,z),up_rr,'numpy')
    func_up_z = sp.lambdify((r,z),up_z,'numpy')
    func_up_zz = sp.lambdify((r,z),up_zz,'numpy')
    func_us = [sp.lambdify((r,z),u,'numpy') for u in us]
    func_us_r = [sp.lambdify((r,z),u,'numpy') for u in us_r]
    func_us_rr = [sp.lambdify((r,z),u,'numpy') for u in us_rr]
    func_us_z = [sp.lambdify((r,z),u,'numpy') for u in us_z]
    func_us_zz = [sp.lambdify((r,z),u,'numpy') for u in us_zz]
    
    
    #This code defines the boundary conditions (see Freidberg for more detail, note sign error in Freidberg for curvature coefficients N1 and N2)
    N1 = -(1+delta)**2/(epsilon*kappa**2)
    N2 = (1-delta)**2/(epsilon*kappa**2)
    N3 = -kappa/(epsilon*np.cos(delta)**2)
    A = np.zeros((7,7))
    b = np.zeros(7)
    A[0] = np.array([func_us[i](1+epsilon,0) for i in range(7)])
    b[0] = -func_up(1+epsilon,0)
    A[1] = np.array(
        [
            func_us_zz[i](1+epsilon,0) + N1 * func_us_r[i](1+epsilon,0) for i in range(7)
        ]
    )
    b[1] = -func_up_zz(1+epsilon,0) - N1 * func_up_r(1+epsilon,0)
    A[2] = np.array([func_us[i](1-epsilon,0) for i in range(7)])
    b[2] = -func_up(1-epsilon,0)
    A[3] = np.array(
        [
            func_us_zz[i](1-epsilon,0) + N2 * func_us_r[i](1-epsilon,0) for i in range(7)
        ]
    )
    b[3] = -func_up_zz(1-epsilon,0) - N2 * func_up_r(1-epsilon,0)
    A[4] = np.array([func_us[i](1-delta*epsilon,kappa*epsilon) for i in range(7)])
    b[4] = -func_up(1-delta*epsilon,kappa*epsilon)
    A[5] = np.array([func_us_r[i](1-delta*epsilon,kappa*epsilon) for i in range(7)])
    b[5] = -func_up_r(1-delta*epsilon,kappa*epsilon)
    A[6] = np.array(
        [
            func_us_rr[i](1-delta*epsilon,kappa*epsilon) + N3 * func_us_z[i](1-delta*epsilon,kappa*epsilon) for i in range(7)
        ]
    )
    b[6] = -func_up_rr(1-delta*epsilon,kappa*epsilon) - N3 * func_up_z(1-delta*epsilon,kappa*epsilon)
    
    #This line solves the boundary conditions to get the coefficients in front of the basis functions for the complementary solution.
    cs=np.linalg.solve(A, b)
    
    #Psi is sum of particular solution and complementary function
    psi = up
    for i in range(7):
        psi+= cs[i]*us[i]
    psi = psi*psi0

    #print(f'\t c0={cs[0]}')    
    #print(f'\t c1={cs[1]}')
    #print(f'\t c2={cs[2]}')
    #print(f'\t c3={cs[3]}')
    #print(f'\t c4={cs[4]}')
    #print(f'\t c5={cs[5]}')
    #print(f'\t c6={cs[6]}')
        
    return r,z,psi





print("Calculations...")


#Solves for analytical solution for psi.
alpha = A/(C-A)
psi0 = C-A #turns out to just be overall normalization factor

r,z,psi = solve_solovev_grad_shafranov(epsilon,kappa,delta,alpha,psi0)

#Defines a numerical grid across the plasma cross-section, just for numerically computing interesting values.
rs = np.linspace(1-epsilon,1+epsilon,500)
zs = np.linspace(-epsilon*kappa,epsilon*kappa,500)
dr = rs[1]-rs[0]
dz = zs[1]-zs[0]
R,Z = np.meshgrid(rs,zs,indexing='ij')
func_psi = sp.lambdify((r,z),psi,'numpy')
Psi = func_psi(R,Z)
R_masked = np.where(Psi<=0,R,None)
Z_masked = np.where(Psi<=0,Z,None)
Psi_masked = np.where(Psi<=0,Psi,None)
R_cleaned = np.array([x for x in np.ravel(R_masked) if x is not None])
Z_cleaned = np.array([x for x in np.ravel(Z_masked) if x is not None])
Psi_cleaned = np.array([x for x in np.ravel(Psi_masked) if x is not None])
# assert external_coil_current**2 + 2*A*Psi_cleaned.any() > 0, 'Toroidal field is vanishing somewhere! This is because plasma poloidal current is cancelling out the effect of the toroidal field coils.'
if external_coil_current**2+2*A*Psi_cleaned.any() < 0:
    print('!!!!ALERT ALERT ALERT!!!!')
    print('\tToroidal current vanishes/changes direction. You have created a field-reversed configuration! Increase external_coil_current to avoid this')
    print('!!!!ALERT ALERT ALERT!!!!')
K3 = dr*dz*np.sum(R_cleaned)
K2 = -dr*dz*np.sum(R_cleaned*Psi_cleaned)
K1 = dr*dz*np.sum((-1*alpha+(1+alpha)*R_cleaned**2)/R_cleaned)
B0 = np.sqrt(external_coil_current**2 + 2*A*func_psi(1,0))
magnetic_axis_position = R_cleaned[np.argmin(Psi_cleaned)]
over_qstar = 1/(2*np.pi) * (2/(1+kappa**2)) * psi0/(epsilon**2 * B0) * K1
beta_p = 8*np.pi**2 * epsilon**2 * (1+alpha) * ((1+kappa**2)/2) * K2/(K1**2 * K3)
beta_t = 8*np.pi**2 * epsilon**4 * (1+alpha) * over_qstar**2 * ((1+kappa**2)/2)**2 *K2/K1**2/K3

#Safety factor on boundary calculation:
tau = np.linspace(0,2*np.pi,200)
dtau = tau[1]-tau[0]
gamma = sp.sqrt(sp.diff(psi,r)**2+sp.diff(psi,z)**2)
func_gamma = sp.lambdify((r,z),gamma,'numpy')
boundary_r = 1 + epsilon * np.cos(tau + np.arcsin(delta)*np.sin(tau))
boundary_z = epsilon * kappa * np.sin(tau)
gamma_on_bound = func_gamma(boundary_r,boundary_z)
dr_dtau = -epsilon * np.sin(tau + np.arcsin(delta) * np.sin(tau))*(1+np.arcsin(delta) * np.cos(tau))
dz_dtau = epsilon*kappa*np.cos(tau)
q_onboundary = external_coil_current/2/np.pi * dtau * np.sum(np.sqrt(dr_dtau**2 + dz_dtau**2)/boundary_r/gamma_on_bound)

print("Calculations...")
print("Calculations...\n")

print('Outputs:')
print(f'\tposition of magnetic axis is R={magnetic_axis_position:.{sig_figs}g}')
print(f'\tpoloidal beta is {beta_p:.{sig_figs}g}')
print(f'\ttoroidal beta is {beta_t:.{sig_figs}g}')
print(f'\tbeta is {1/(1/beta_p + 1/beta_t):.{sig_figs}g}')
print(f'\tsafety factor on boundary is {q_onboundary:.{sig_figs}g}')
print(f'\tkink safety factor is {1/over_qstar:.{sig_figs}g}')
print(f'\ttoroidal beta times kink safety factor squared on boundary is {beta_t/over_qstar**2:.{sig_figs}g}')



#PLOTTING
rs = np.linspace(0.8*(1-epsilon),1.2*(1+epsilon),1000)
zs = np.linspace(-1.2*(epsilon*kappa),1.2*(epsilon*kappa),1000)
R,Z = np.meshgrid(rs,zs,indexing='ij')
Psi = func_psi(R,Z)
fig,ax = plt.subplots(figsize=(7,7))
thing=ax.contourf(R,Z,Psi,levels=np.linspace(np.min(Psi),-np.min(Psi),200),cmap='RdBu')
ax.contour(R,Z,Psi,levels=np.linspace(np.min(Psi),0,12),colors='gray',linestyles='solid')
tau = np.linspace(0,2*np.pi,200)
boundary_r = 1 + epsilon * np.cos(tau + np.arcsin(delta)*np.sin(tau))
boundary_z = epsilon * kappa * np.sin(tau)
ax.contour(R,Z,Psi,levels=np.linspace(0,-np.min(Psi),12),colors='gray',linestyles='dashed')
ax.plot(boundary_r,boundary_z,color='green',label='boundary')
ax.legend()
fig.colorbar(mappable=thing)
ax.set_ylabel('Z')
ax.set_xlabel('R') 
ax.set_aspect('equal')  
ax.set_title(r'$\psi(R,Z)$')
plt.savefig('flux_surface_cross_section.png', dpi=300)

tau = np.linspace(0,2*np.pi,50)
theta = np.linspace(0, 2*np.pi, 100) 
tau, THETA = np.meshgrid(tau, theta)
boundary_R = 1 + epsilon * np.cos(tau + np.arcsin(delta)*np.sin(tau))
boundary_Z = epsilon * kappa * np.sin(tau)
X = boundary_R * np.cos(THETA)
Y = boundary_R * np.sin(THETA)
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X, Y, boundary_Z, cmap='viridis', alpha=0.8)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
ax.set_box_aspect([1, 1, 1])  
ax.set_aspect('equal')  
plt.savefig('boundary_surface_3D.png', dpi=300)
print('Plots saved.')