# Name: IC.py
# Purpose: Calculate the Ikeda-Carpenter function.
#
# Date: October 15th, 2018
# Author: Timothy R. Prisk
# Contact: timothy.prisk@nist.gov
#
# Notes: This program calculates the velocity-time structure of a neutron pulse
# at the ARCS spectrometer.
#

# Import useful packages.
import numpy as np
import scipy
from scipy import special
from pylab import *
import scipy.special as sps

# Constants
alpha_s = 1.0 #1/us
beta_s = 0.15 #1/us
R_s = 0.25 #dimensionless
sigma_i = 3.04 #us
L_fermi = 11.61 #meters
L_m1 = 11.835 #meters
L_m2 = 18.50 #meters

def IC(alpha, beta, R, t):
    return 0.5*alpha*((1.0 - R)*alpha**2*t**2*np.exp(-alpha*t) +
    ((2.0*R*alpha**2*beta)/(alpha - beta)**3)* (np.exp(-beta*t)
    - np.exp(-alpha*t)*(1.0 + (alpha - beta)*t + 0.5*(alpha-beta)**2*t**2)))

def meantime(alpha, beta, R):
    return 3.0/alpha + R/beta

def Gaussian(sigma, t):
    return (1.0/np.sqrt(2.0 * np.pi))*(1.0/sigma)*np.exp(-0.5*t**2/sigma**2)

def ICg(alpha, beta, R, sigma, t):
    prefactor = alpha/(2.0 * np.sqrt(2.0*np.pi) * sigma)
    exp_a = np.exp(-alpha*t)
    c1 = (1.0 - R)*alpha**2
    d = alpha - beta
    c2 = 2.0*R*alpha**2*beta/d**3

    prefactor_gn = np.sqrt(np.pi/2.0)*sigma

    ga0 = prefactor_gn*np.exp(0.5*alpha**2*sigma**2)*sps.erfc((alpha*sigma -
    t/sigma)/np.sqrt(2.0))
    gb0 = prefactor_gn*np.exp(0.5*beta**2*sigma**2)*sps.erfc((beta*sigma -
    t/sigma)/np.sqrt(2.0))

    ga1 = sigma**2*(-alpha*ga0 + np.exp(alpha*t - t**2/2.0/sigma**2))
    ga2 = sigma**2*(ga0 - alpha*ga1 - t*np.exp(alpha*t - t**2/2.0/sigma**2))

    term1 = (c1*t**2 - c2*(1.0 + d*t + 0.5*d**2*t**2))*exp_a*ga0
    term2 = c2*np.exp(-beta*t)*gb0
    term3 = (2*c1*t - c2*d*(1.0+d*t))*exp_a*ga1
    term4 = (c1 - c2*d**2/2.0)*exp_a*ga2

    return prefactor*(term1 + term2 + term3 + term4)

def vi(Ei):
    return np.sqrt(Ei/5.227) * 1./0e-3 #from km/s to m/us

def sigma_det(Ei):
    Lmon = 0.01 #m thickness of the monitor
    return np.sqrt(((Lmon/vi(Ei)) + 1.0**2  )/12.0)


def sigma_m2(sigma_m1):
    return (L_m2/L_m1) * sigma_m1

def sigma_m2_corrected(sigma_m1, Ei):
    return np.sqrt(sigma_det**2 + sigma_m2**2)


def alpha_m2(alpha):
    return (L_fermi / (L_m2 - L_fermi)) * alpha

def beta_m2(beta):
    return (L_fermi / (L_m2 - L_fermi)) * beta


#One to fit with.
def ICg_fit(scale, tposition, a, b, R, s, t):
    return scale*ICg(a, b, R, s, tposition - t)

def alpha_source(a):
    return a * (L_m2 - L_fermi)/L_fermi

def beta_source(b):
    return b * (L_m2 - L_fermi)/L_fermi





ts = np.arange(0.0, 1000.0, 0.5)
ICs = np.zeros(len(ts))
IC2s = np.zeros(len(ts))
ICgs = np.zeros(len(ts))
mon2 = np.zeros(len(ts))

for i in range(len(ts)):
    ICs[i] = IC(alpha_s, beta_s, R_s, ts[i])
    IC2s[i] = IC(0.25, 0.05, 0.8, ts[i])
    ICgs[i] = ICg(alpha_s, beta_s, R_s, sigma_s, ts[i])
    mon2[i] = ICg(alpha_m2(alpha_s), beta_m2(beta_s), R_s, sigma_m2(sigma_i),
    500.0 - ts[i])

plt.plot(ts, ICs, 'ko')
#plt.plot(ts, IC2s, 'bo')
plt.plot(ts, ICgs, 'ro')
plt.plot(ts, mon2, 'yo')
plt.show()
