#Code to compute the power spectrum from the solution to the Mukhanov equation. 
import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
import math 
plt.rcParams["font.size"] = 6
plt.autoscale(enable=True, axis='both', tight=None)

#Define Parameters.
n = float(input('Enter the value of n: '))
k_a = np.geomspace(start = 0.005, stop = 10, num = 1000)
t = np.linspace(start = 60.0, stop = 0.0, num = 1000) #t is the efold number.
k_0 = 0.005
k_max = 1e16
#########################################
#Slow-roll parameters (using phi^2 for now and t is the e-fold number)
 
def ep_func(t, n):
    return (n**2)/(4*n*t + n**2)

def H_func(t, n):
    return ((2*n*t + (n**2)/2)**(n/2))/3

def eta_func(t, n):
    return 2*((2*n - n**2)/(2*n*t + ((n**2)/2)))

def a_func(t):
    return np.exp(t)
##########################################
#n_s and r

def r_func(t, n):
    return (16*(n**2))/(4*n*t + n**2)

def n_func(t, n):
    return 1 -(6*(n**2)/(4*n*t + n**2)) + (2*(n*(n-1)))/(2*n*t + (n**2)/2)
  
##########################################
#Create the two simplified ODE's.
#This is in terms of the e-fold number not confromal time. 
def sol(Y, t, k, ep, eta, H, a):
    return np.array([Y[1], -(3 - ep_func(t, n) + eta_func(t, n))*Y[1] -((k**2)/((a_func(t)**2)*H_func(t, n)**2))*Y[0]])
##########################################
# Find the solution using Scipy odeint function while defining the I.C's.
Yr_0 = [1/(np.sqrt(2*k_max)), 0]
Yi_0 = [0, -np.sqrt(k_max/2)]
solutions = []
# Real part of the Mukhanov equation.
for k in k_a:
    asolr = odeint(sol, Yr_0, t, args=(k, ep_func(t, n), eta_func(t, n), H_func(t, n), a_func(t)))    
    solutions.append(asolr)

# Imaginary Part of the Mukhanov equation.
for k in k_a:
    asoli = odeint(sol, Yi_0, t, args=(k, ep_func(t, n), eta_func(t, n), H_func(t, n), a_func(t)))    
    solutions.append(asoli)

Yr = asolr[:,0]
Yi = asoli[:,0]
########################################## 
#Plot the solution to the mukhanov equation. 
def mukh(Yi):
    return np.absolute(Yi)
m = mukh(Yi)
plt.plot(t, m, 'r-')
plt.yscale('log') 
##########################################
#Define the the power spectrum.

#Mukhanov Power spectrum
def m_spec(k_a, Yr, Yi):
    return 2.1e-9*((k_a**3)/0.005)*(np.absolute(Yr)**2 + np.absolute(Yi)**2)

#Slow-roll power spectrum 
def s_spec(k_a, t, n):
    return 2.1e-9*(k_a/0.005)**(-(2*n**2 + n)/(2*n*t + (n**2)/2))
 
pm = m_spec(k_a, Yr, Yi)
ps = s_spec(k_a, t, n)

#Plot the power spectrums and parameters in subplots
fig = plt.figure()

plt.subplot(2, 2, 1)
plt.plot(k_a, pm, 'r-')
plt.title('Power spectrum P(k) for n = ' +str(n))
plt.xlabel('wavenumber (k)')
plt.ylabel('P(K)')
plt.yscale('log')
plt.xscale('log')

plt.subplot(2, 2, 2)
plt.plot(k_a, ps, 'b-')
plt.title('Slow-roll Power spectrum P(k) for n = ' +str(n))
plt.xlabel('wavenumber (k)')
plt.ylabel('P(K)')
plt.yscale('log')
plt.xscale('log')

plt.subplot(2, 2, 3)
plt.plot(t, H_func(t, n), 'y-')
plt.title('Epsilon Vs e-fold')
plt.xlabel('e-fold')
plt.ylabel('$H$(N)')
#plt.yscale('log')

plt.subplot(2, 2, 4)
plt.plot(t, n_func(t, n), 'g-')
plt.title('$n_s$ Vs e-fold')
plt.xlabel('e-fold')
plt.ylabel('$n_s$')
#plt.yscale('log')

plt.tight_layout()
plt.show()