#!/bin/sh  
# \
exec oagtclsh "$0" "$@"

# Computation of elastic gas scattering lifetime using DA and twiss output from elegant, along with
# s-dependent pressure profiles for various gases.
# The known gases are H2, H2O, CH4, CO, CO2, and N2. If columns with these names are in the 
# pressure input file, they are used in the computations. Units must be nT or nTorr.
# M. Borland, 2014.

if {![info exists env(OAG_TOP_DIR)]} { set env(OAG_TOP_DIR) /usr/local }
set auto_path [linsert $auto_path 0  $env(OAG_TOP_DIR)/oag/apps/lib/$env(HOST_ARCH)]

APSStandardSetup

set usage {usage: elasticScatteringLifetimeDetailed -aperture <filename> [-name <string>] -twiss <filename> -pressure <filename> -output <filename> [-sPeriodic meters(27.6)]}
set aperture ""
set twiss ""
set name ""
set pressure ""
set output ""
set sPeriodic 27.6
set args $argv
if {[APSStrictParseArguments {aperture twiss name pressure output sPeriodic}] || ![string length $aperture] || ![string length $twiss] || ![string length $pressure] || \
        ![string length $output] || [expr $sPeriodic<=0]} {
    return -code error "$usage"
}
foreach input {aperture twiss pressure} {
    if ![file exists [set $input]] {
        return -code error "not found: [set $input]"
    }
}
if [file exists $output] {
    return -code error "in use: $output"
}

# Set up arrays for gases
set C(H2) 1
set Z(H2,1) 1
set S(H2,1) 2

set C(H2O) 2
set Z(H2O,1) 1
set S(H2O,1) 2
set Z(H2O,2) 6
set S(H2O,2) 1

set C(CH4) 2
set Z(CH4,1) 8
set S(CH4,1) 1
set Z(CH4,2) 1
set S(CH4,2) 4

set C(CO) 2
set Z(CO,1) 8
set S(CO,1) 1
set Z(CO,2) 6
set S(CO,2) 1

set C(CO2) 2
set Z(CO2,1) 8
set S(CO2,1) 1
set Z(CO2,2) 6
set S(CO2,2) 2

set C(N2) 1
set Z(N2,1) 7
set S(N2,1) 2

set knownSpeciesList [array names C]

# Determine what species are in the pressure file
set columnList [exec sddsquery -column $pressure]
set speciesList ""
foreach column $columnList {
    if [lsearch -exact $knownSpeciesList $column]!=-1 {
        lappend speciesList $column
    }
}
if [llength $speciesList]==0 {
    return -code error "no known gas species ([join $knownSpeciesList ,]) identified in list of columns: [join $columnList ,]"
}

# Check units of pressure values
array set units [exec sddsquery -column $pressure -appendunits=bare]
foreach species $speciesList {
    if [lsearch -exact [list nT nTorr] $units($species)]==-1 {
        return -code error "units for $species are not valid: $units($species); use nT or nTorr"
    }
}

# Compute Fg factor for each species
# Fg = sum(a=1,Cg) E(g,a)*S(g,a), where E(g,a) = (Z(g,a)*re/(2*gamma))^2
set gamma [lindex [exec sdds2stream -parameter=pCentral $twiss] 0]
set re 2.81794092e-15
set PIntegList ""
set PConv [expr 133.3224*1e-9]
set rateSumOpt "-define=column,rate,0 "
set T [expr 20+273.15]
foreach g $speciesList {
    set Fg 0.0
    for {set a 1} {$a<=$C($g)} {incr a} {
        set Fg [expr $Fg+$Z($g,$a)*$Z($g,$a)*$S($g,$a)]
    }
    set Fg [expr $Fg*pow($re/(2*$gamma), 2)]
    # ng(s) is the number density
    lappend PIntegList "-define=column,ng-$g,$g $PConv * kb_mks / $T /,units=1/m^3" 
    # Fg*Ig(s) = Fg*ng(s)*Ie(s) where Ie(s) is the integral over the acceptance for elastic scattering
    lappend PIntegList "-define=column,FgIg-$g,ng-$g $Fg * CSInteg *"
    lappend PIntegList "-define=column,localRate-$g,FgIg-$g c_mks * sMax /,units=1/s/m" 
    lappend PIntegList "-process=localRate-$g,integ,rate-$g,functionOf=s"
    # Rate is the average value of Fg*Ig(s) times c
    append rateSumOpt "rate-$g + "
}

set tmpRoot [APSTmpString]
APSAddToTempFileList $tmpRoot.1 $tmpRoot.2 $tmpRoot.3

# Find beta functions at ID straight, use only first sector's worth of data
# Sort into increasing s order, remove duplicates
exec sddsprocess $twiss -pipe=out -process=beta*,first,%s0 -filter=col,s,0,[expr $sPeriodic+1e-3] \
    | sddssort -pipe -column=s -unique \
    | sddsconvert -pipe=in $tmpRoot.1 -retain=col,betax,betay,s -retain=parameter,beta?0

# Interpolate the pressure data at locations of lattice function data 
# Convert all the s-dependent quantities into parameters (one s location per page)
exec sddsinterp $pressure -pipe=out -fileValues=$tmpRoot.1,column=s -column=s,[join $speciesList ,] -belowRange=abort -aboveRange=abort \
    | sddsxref -pipe $tmpRoot.1 -take=* -transfer=param,* \
    | sddsexpand -pipe=in $tmpRoot.2 
set nTw [exec sdds2stream -npage=bare $tmpRoot.2]

# Replicate the twiss and pressure data once for each DA contour in the aperture file
set nAp   [exec sdds2stream -npage=bare $aperture]
if $nAp>1 {
    eval exec sddscombine [APSReplicateItem -item $tmpRoot.2 -number $nAp] $tmpRoot.3
    set twissFile $tmpRoot.3
} else {
    set twissFile $tmpRoot.2
}


# 1. Replicate the aperture data once for each point in the twiss data file
# 2. Pull in the twiss and pressure data as parameters (now each page is for a specific s)
# 3. Perform integral over phi for each s.
# 4. Perform integral over s for each species.
# 5. Sum over species to get total rate.

eval exec sddscombine [APSReplicateItem -number $nTw -item $aperture] -pipe=out -retain=col,*$name \
    | sddsxref -pipe $twissFile -leave=* -transfer=param,* \
    | sddsprocess -pipe \
    "{-define=col,phi,x$name y$name atan2}" \
    "{-define=col,theta,x$name sqr betax / betax0 / y$name sqr betay / betay0 / + sqrt}" \
    "{-define=col,CS,4 theta 2 / tan sqr /}" \
    | sddssort -pipe -column=phi,incr -unique \
    | sddsprocess -pipe \
    -process=CS,integral,%sInteg,functionOf=phi \
    | sddscollapse -pipe \
    | sddsbreak -pipe -decreaseOf=s \
    | sddsprocess -pipe \
    -process=s,max,%sMax \
    $PIntegList \
    | tee $output.detail \
    | sddscollapse -pipe \
    | sddsprocess -pipe=in $output \
    "{$rateSumOpt}" \
    "{-define=column,elasticScatteringLifetime,rate rec 3600 /,units=h}" \
    -process=elasticScatteringLifetime,ave,%sAve \
    -process=elasticScatteringLifetime,stand,%sStDev 



