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

set auto_path [linsert $auto_path 0 /usr/local/oag/apps/lib/$env(HOST_ARCH)]
set auto_path [linsert $auto_path 0 /usr/local/oag/lib_patch/$env(HOST_ARCH)]

APSStandardSetup

set usage "optimizeKickersSmartly -logDir <dirname> -average <number> -stored <0|1> -mode <X|Y|XY> \[-doSetup <0|1>\] \[-fraction <0.75>\] \[-nIterations <number>\] \[-verbose 1\]"

set args $argv
set logDir ""
set average ""
set doSetup 0
set fraction 0.75
set mode ""
set stored ""
set nIterations 5
set verbose 0
APSParseArguments {logDir average doSetup fraction mode verbose stored nIterations}
if {![string length $logDir] || ![string length $average] || ![string length $mode] || ![string length $stored]} {
    puts stdout $usage
    exit
}
set simulRun 0
#set execPath /home/oxygen/SAJAEV/APPS/FPGA-Bpms
set execPath /home/helios/SR/bin

#---------------------------------------------------------------------------------------
proc SetupTrigger {} {
    global fpgaSectorList
    set selectedIOCList ""
    foreach sector $fpgaSectorList {lappend selectedIOCList S${sector}A S${sector}B}
    set postTrigger 260144
    #--- Set postrigger for all IOCs
    if [catch {exec cavput -list=[join $selectedIOCList ,] -list=:turn:gtr:numberPTS=$postTrigger -pend=10} result] {
	return -code error "Error setting up trigger(4): $result" 
    }
    #--- Set trigger to event and arm it
    if [catch {exec cavput -list=[join $selectedIOCList ,] \
		   -list=:turn:gtr:trigger=event,:turn:gtr:arm=1 -pend=10} result] {
        return -code error "Error setting up trigger(1): $result"
    }
    #--- Disable/enable required triggers
    if [catch {exec cavput -list=[join $selectedIOCList ,] \
		   -list=:clientEvent2d.OUT4=Disabled,:clientEvent2d.OUT5=Enabled,:clientEvent33.OUT4=Disabled \
		   -pend=10} result] {
        return -code error "Error setting up trigger(2): $result"
    }
}
#---------------------------------------------------------------------------------------
proc WaitForTriggers {args} {
    global fpgaSectorList
    set selectedIOCList ""
    foreach sector $fpgaSectorList {lappend selectedIOCList S${sector}A S${sector}B}
    set wait 1
    set counter 0
    while {$wait == 1} {
	set triggerList [join [exec cavget -list=[join $selectedIOCList ,] -list=:TurnHistoryIsArmed] ]
	if {[lsearch -exact $triggerList 1] == -1} {set wait 0}
	after 100
	incr counter
	if {$counter > 50} {return -code error "Waiting for triggers to discharge takes too long..."}
    }
}
#---------------------------------------------------------------------------------------

proc CollectInjectedMotion {args} {
    set stored 0
    APSParseArguments {wmonFile smonFile outputFile nTrajAver plane simulRun stored verbose}
    if !$simulRun {
	if [file exists $outputFile] {
	    return -code error "Error: File $outputFile already exists."
	}
    }
    #--- Setup injection cycle and bucket:
    if !$simulRun {exec cavput -list=Mt:SRbunch000.VAL=0,Mt:SRinjectNumBunchesAO.VAL=1,Mt:SRinjectMaxCyclesAO.VAL=1 -pend=10}
    if $stored {
	if $verbose {puts stdout "Disabling injectors..."}
	#--- Disable injectors:
	if [catch {SetPulsedMagnetEnables -location PAR -state 0} result] {
	    puts stderr "SetPulsedMagnetEnables: $result"; exit 1
	}
	#--- Turn off IS1 discharge (don't want the stray fields in the beam way) and disable RTFB freeze:
	if [catch {exec cavput -list=Mt:Ddg3chan2.GATE=0,SRFB:GBL:DisableFreezeHBO.VAL=1,SRFB:GBL:DisableFreezeVBO.VAL=1 \
		       -pend=5} result] {
	    after 1000
	    exec cavput -list=Mt:Ddg3chan2.GATE=0,SRFB:GBL:DisableFreezeHBO.VAL=1,SRFB:GBL:DisableFreezeVBO.VAL=1 -pend=5
	}
    }
    #--- Collect nTrajAver measurements:
    set fileList ""
    for {set i 0} {$i < $nTrajAver} {incr i} {
	#--- Setup triggers:
	if $simulRun {
	    puts stdout "Setup triggers..."
	} else {
	    if $verbose {puts stdout "Setting up triggers..."}
	    if [catch {SetupTrigger} result] {
		return -code error "SetupTrigger: $result"
	    }
	}
	#--- Press "Inject" button:
	if $simulRun {
	    puts stdout "cavput -list=Mt:SRinjectMultiBO.VAL=1"
	} else {
	    if $verbose {puts stdout "Pressing \"inject\" button..."}
	    exec cavput -list=Mt:SRinjectMultiBO.VAL=1
	}
	#--- Capture data:
	if !$simulRun {
	    if $verbose {puts stdout "Waiting for trigger to discharge..."}
	    if [catch {WaitForTriggers} result] {
		after 1000
		if [catch {WaitForTriggers} result] {
		    return -code error "WaitForTriggers: $result"
		}
	    }
	    if $verbose {puts stdout "Collecting waveforms..."}
	    if [catch {exec sddswmonitor $wmonFile $outputFile.$i -steps=1 -scalars=$smonFile} result] {
		return -code error "sddswmonitor: $result"
	    }
	}
	if [catch {ProcessInjectionTrajectory -dataFile $outputFile.$i -plane $plane} result] {
	    puts stdout "Error running \"ProcessInjectionTrajectory -dataFile $outputFile.$i -plane $plane\""
	    puts stdout "ProcessInjectionTrajectory: $result"
	    puts stdout "Continue measurements..."
	}
	lappend fileList $outputFile.$i
    }
    #--- Enable IS1 discharge and RTFB freeze back:
    if !$stored {
	exec cavput -list=Mt:Ddg3chan2.GATE=1,SRFB:GBL:DisableFreezeHBO.VAL=0,SRFB:GBL:DisableFreezeVBO.VAL=0 -pend=5
    }

    #--- Combine files and average:
    if {[llength $fileList] == 1} {
	exec sddsprocess $fileList $outputFile -reprint=para,Plane,$plane -redef=para,NTrajectories,$nTrajAver,type=long
    } else {
	eval exec sddscombine $fileList -pipe=out \
	    | sddsenvelope -pipe -copy=Index -mean=*position,*sum -stand=*position,*sum \
	    | sddsconvert -pipe -edit=col,*Mean,%@Mean@@ \
	    | sddsxref -pipe [lindex $fileList 0] -leave=* -transfer=par,* \
	    | sddsprocess -pipe=in $outputFile -reprint=para,Plane,$plane -redef=para,NTrajectories,$nTrajAver,type=long
    }
    #--- Creating trajectory and oscillation files:
    if [catch {ProcessInjectionTrajectory -dataFile $outputFile -plane $plane} result] {
	return -code error "ProcessInjectionTrajectory: $result"
    }
}

#---------------------------------------------------------------------------------------

proc ProcessInjectionTrajectory {args} {
    global fpgaSectorList timingShiftSector turn0 sector39L sectorL OAGGlobal
    APSParseArguments {dataFile plane}
    set twiFile $OAGGlobal(SRLatticesDirectory)/default/aps.twi
    
    #--- Subtract DC orbit:
    exec sddsprocess $dataFile $dataFile.oscil -process=S*position,aver,%sAver,functionof=Index,lower=1000,upper=2000 \
	"-redef=col,%s,%s %sAver -,select=S*position"
    if {[string compare $plane "XY"] == 0} {
	set planeList [list X Y]
	set bpmSufListList [list [list A:P2 B:P5 B:P2] [list A:P3 A:P4 B:P4 B:P3]]
    } else {
	set planeList $plane
	set bpmSufListList [list [list A:P2 A:P3 A:P4 B:P5 B:P4 B:P3 B:P2]]
    }
    foreach localPlane $planeList bpmSufList $bpmSufListList {
	#--- Determine if the dataFile has StDev columns and defining suffix lists:
	set inputColumnList [join [exec sddsquery $dataFile -col] ]
	if {[lsearch -regexp $inputColumnList StDev] != -1} {
	    set suffixList [list position positionStDev sum sumStDev]
	    set pvSuffixList [list :turnHistorySmall:${localPlane}position :turnHistorySmall:${localPlane}positionStDev \
				  :turnHistorySmall:${localPlane}sum :turnHistorySmall:${localPlane}sumStDev]
	    set columnList [list Trajectory TrajectoryStDev Sum SumStDev]
	} else {
	    set suffixList [list position sum]
	    set pvSuffixList [list :turnHistorySmall:${localPlane}position :turnHistorySmall:${localPlane}sum]
	    set columnList [list Trajectory Sum]
	}
	#--- Processing file:
	set fileList ""
	foreach suffix0 $suffixList pvSuffix $pvSuffixList columnName $columnList {
	    set suffix ${localPlane}$suffix0
	    set columnList1 ""
	    set columnList2 ""
	    foreach sector $fpgaSectorList {
		foreach bpm $bpmSufList {
		    if {$sector < $timingShiftSector || $sector == 40} {
			lappend columnList1 S${sector}${bpm}:turnHistorySmall:$suffix
		    } else {
			lappend columnList2 S${sector}${bpm}:turnHistorySmall:$suffix
		    }
		}
	    }
	    exec sddsconvert $dataFile.oscil -pipe=out -retain=col,Index,[join $columnList1 ,] \
		| sddsprocess -pipe=in $dataFile.oscil.1 "-redef=col,Index0,Index $turn0 -" \
		-filter=col,Index0,-10,100
	    exec sddsconvert $dataFile.oscil -pipe=out -retain=col,Index,[join $columnList2 ,] \
		| sddsprocess -pipe=in $dataFile.oscil.2 "-redef=col,Index0,Index $turn0 - 1 -" \
		-filter=col,Index0,-10,100
	    exec sddsxref $dataFile.oscil.1 $dataFile.oscil.2 -pipe=out -take=S* \
		| sddsbreak -pipe -rowlimit=1 \
		| sddsprocess -pipe -process=Index0,first,TurnAfterInj \
		| sddsconvert -pipe -del=col,Index,Index0 \
		| sddstranspose -pipe -oldColumn=BPMName -rootname=$columnName \
		| sddsprocess -pipe "-reedit=col,BPMName,%@$pvSuffix@@" -reprint=para,Plane,$localPlane \
		| sddsxref -pipe $twiFile -take=s -match=BPMName=ElementName -nowarning -reuse=page \
		| sddsprocess -pipe "-redef=col,s,s $sector39L > ? s $sector39L - : s $sectorL + \$" \
		| sddssort -pipe=in $dataFile.$suffix0 -col=s
	    file delete $dataFile.oscil.1 $dataFile.oscil.2
	    lappend fileList $dataFile.$suffix0
	}
	#--- Combining all columns:
	eval exec sddsxref $fileList $dataFile.$localPlane.trajectory -take=* -leave=BPMName,s
	eval file delete $fileList
    }
}

#---------------------------------------------------------------------------------------

#--- Procedure from FPGABpmTurnHistory by H. Shang, adapted
proc LoadRAM {args} {
    global fpgaSectorList logDir scalarList plane
    APSParseArguments {plane mode verbose}
    set SCRdir /home/helios/oagData/SCR/snapshots/SBPMWaveform
    set logFile /tmp/[APSTmpString]
    switch -exact $mode {
	singleBunch {
	    set refFile SBPMWaveform-singleBunch.gz
	}
	24singlets {
	    set refFile SBPMWaveform-24singlets.gz
	}
	measurement {
	    switch -exact $plane {
		X {set refFile SBPMWaveform-singleX.gz}
		Y {set refFile SBPMWaveform-singleY.gz}
		XY {set refFile SBPMWaveform-singleXYadt.gz}
	    }
	}
    }
    #--- Make monitor files if mode=measurement:
    if {[string compare $mode measurement] == 0} {
	#--- Make wmonitor files:
	if {[string compare $plane "XY"] == 0} {
	    set bpmPlaneList [list X Y Y X Y Y X]
	} else {
	    set bpmPlaneList [list $plane $plane $plane $plane $plane $plane $plane]
	}
	foreach sector $fpgaSectorList {
	    foreach bpm [list A:P2 A:P3 A:P4 B:P5 B:P4 B:P3 B:P2] bpmPlane $bpmPlaneList {
		lappend columnList S${sector}${bpm}:turnHistorySmall:${bpmPlane}position
		lappend columnList S${sector}${bpm}:turnHistorySmall:${bpmPlane}sum
	    }
	}
	set wmonFile $logDir/fpga.wmon
	set smonFile $logDir/fpga.mon
	exec sddsmakedataset $wmonFile \
	    -para=WaveformLength,type=long -data=4096 \
	    -col=WaveformName,type=string -data=[join $columnList ,] \
	    -col=WaveformPV,type=string -data=[join $columnList ,]
	exec sddsmakedataset $smonFile \
	    -col=ControlName,type=string -data=[join $scalarList ,]
    }
    #---
    set filename [file readlink $SCRdir/$refFile]
    set restoreError 0
    if [catch {APSSCRLogAction -file $refFile -action restore-start} result] {
	puts stdout "$result"
    }
    set configDescription [exec sdds2stream $filename -para=SnapshotDescription]
    if $verbose {puts stdout "Restoring BPM configuration: $configDescription"}
    if [catch {exec sddscasr $filename -restore -logFile=$logFile -pendIOTime=100 \
		   -waveform=directory=$SCRdir,onefile,rootname=[file rootname [file tail $filename]],extension=.waveform.gz } result] {
	puts stdout "Error in loading $refFile to iocs: $result"
	set restoreError 1
    }
    set fd [open $logFile "r"]
    while {![eof $fd]} {
	set text [gets $fd]
	if {[string match -nocase "*error*" $text]} {
	    set restoreError 1
	    break
	}
    }
    close $fd
    if $restoreError {
	APSFileDisplayWindow [APSUniqueName .fileDisp] -fileName $logFile \
	    -comment "Error output from burtwb" -deleteOnClose 1 \
	    -printCommand "enscript -r"
    } else {
	puts stdout "FPGA setup done."
    }
}

#---------------------------------------------------------------------------------------

set fpgaSectorList [list 40 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 29 30 31 32 33 34 35 36 37 38 39]
set bpmList [list A:P2 A:P3 A:P4 B:P5 B:P4 B:P3 B:P2]
set scalarList [list S:IK1:VoltageSetSendAO S:IK2:VoltageSetSendAO \
		    S:IK3:VoltageSetSendAO S:IK4:VoltageSetSendAO S:IS1:VoltageSetSendAO S:IS2:VoltageSetSendAO \
		    BTS:BV1:CurrentAO BTS:CV1:CurrentAO]
set sector39L 1076.4
set sectorL 27.6
set circum 1104
set timingShiftSector 27
set turn0 2096

#--- Traj files start with S40
set sectorStart 1
set sectorEnd 2
set sStart [expr $sectorStart * 27.6]
set sEnd [expr ($sectorEnd + 1) * 27.6]

set wmonFile $logDir/fpga.wmon
set smonFile $logDir/fpga.mon
set outputFileRoot onAxisSmart
set fitRoot smartFit
set fileIndexString 0
set outputFile $logDir/$outputFileRoot$fileIndexString
if {[string first XY $mode] != -1} {
    set plane XY
} else {
    if {[string first X $mode] != -1} {set plane X} else {set plane Y}
}
#--- Press "Gun to Booster Extraction" button:
if [catch {TogglePulsedMagnetEnables -location GuntoBoosterExt} result] {
    return -code error "TogglePulsedMagnetEnables: $result"
}
after 1000

#--- Set up FPGA waveforms:
if $doSetup {
    if [catch {LoadRAM -plane $plane -mode measurement -verbose $verbose} result] {
	puts stdout "LoadRAM: $result"
	exit
    }
}

#--- Collect motion:
catch {eval file delete [glob $logDir/$outputFileRoot*]}
if $verbose {puts stdout "Collecting motion data..."}
if [catch {CollectInjectedMotion -simulRun $simulRun -wmonFile $wmonFile -smonFile $smonFile \
	       -outputFile $outputFile -nTrajAver $average -plane $plane -stored $stored -verbose $verbose} result] {
    puts stdout "CollectInjectedMotion: $result"
}

set outputString "Motion rms: "
foreach localPlane [list X Y] {
    if {[string first $localPlane $plane] != -1} {
	if [catch {exec sddsprocess $outputFile.$localPlane.trajectory -pipe=out -filter=para,TurnAfterInj,0,0 \
		       -filter=col,s,$sStart,$sEnd -process=Trajectory,stand,Std \
		       | sdds2stream -pipe=in -para=Std} motionStDev$localPlane] {
	    return -code error "Error processing $outputFile.$localPlane.trajectory: [set motionStDev$localPlane]"
	}
	append outputString "$localPlane: [set motionStDev$localPlane]  "
    }
}
puts stdout $outputString

set continue 1
set i 0
set iterations $nIterations
while {$continue == 1} {
    incr i
    set oldValues [exec cavget -list=S:IK3:VoltageSetSendAO,S:IK4:VoltageSetSendAO,BTS:BV1:CurrentAO,BTS:CV1:CurrentAO]
    set oldIK3 [lindex $oldValues 0]
    set oldIK4 [lindex $oldValues 1]
    set oldBV1 [lindex $oldValues 2]
    set oldCV1 [lindex $oldValues 3]
    #--- Process motion:
    set localPlane X
    if {[string first $localPlane $plane] != -1} {
	set command "$execPath/processInjectionTrajectory -inputFile $outputFile.$localPlane.trajectory \
		   -outputRoot $fitRoot -plot 0 -verbose 1 -matchX \"0 0\" -sectorStart $sectorStart -sectorEnd $sectorEnd"
	if $verbose {puts stdout $command}
	if [catch {eval exec $command >&@ stdout} result] {
	    return -code error "$execPath/processInjectionTrajectory: $result"
	}
	set ik3 [exec sddsprocess $fitRoot.results -pipe=out -match=col,ValueName=newIK3 \
		     | sdds2stream -pipe=in -col=Value]
	set ik4 [exec sddsprocess $fitRoot.results -pipe=out -match=col,ValueName=newIK4 \
		     | sdds2stream -pipe=in -col=Value]
	set newIK3 [expr $oldIK3 + ($ik3 - $oldIK3) * $fraction]
	set newIK4 [expr $oldIK4 + ($ik4 - $oldIK4) * $fraction]
	set changeIK3 [expr abs($newIK3 - $oldIK3)/$oldIK3]
	set changeIK4 [expr abs($newIK4 - $oldIK4)/$oldIK4]
	if [expr sqrt($changeIK3*$changeIK3 + $changeIK4*$changeIK4) < 0.01] {set continue 0}
	if {$i > $iterations} {set continue 0}
	puts stdout "Iteration $i: Setting new kickers: $newIK3 and $newIK4"
	exec cavput -list=S:IK3:VoltageSetSendAO=$newIK3,S:IK4:VoltageSetSendAO=$newIK4 -pend=5
    }
    set localPlane Y
    if {[string first $localPlane $plane] != -1} {
	set command "$execPath/processInjectionTrajectory -inputFile $outputFile.$localPlane.trajectory \
		   -outputRoot $fitRoot -plot 0 -verbose 1 -matchY \"0 0\""
	if $verbose {puts stdout $command}
	if [catch {eval exec $command >&@ stdout} result] {
	    return -code error "$execPath/processInjectionTrajectory: $result"
	}
	set bv1 [exec sddsprocess $fitRoot.results -pipe=out -match=col,ValueName=deltaBV1 \
		     | sdds2stream -pipe=in -col=Value]
	set cv1 [exec sddsprocess $fitRoot.results -pipe=out -match=col,ValueName=deltaCV1 \
		     | sdds2stream -pipe=in -col=Value]
	set newBV1 [expr $oldBV1 + $bv1 * $fraction]
	set newCV1 [expr $oldCV1 + $cv1 * $fraction]
	if [expr sqrt($newBV1*$newBV1 + $newCV1*$newCV1) < 0.05] {set continue 0}
	if {$i > $iterations} {set continue 0}
	puts stdout "Iteration $i: Setting new correctors: $newBV1 and $newCV1"
	exec cavput -list=BTS:BV1:CurrentAO=$newBV1,BTS:CV1:CurrentAO=$newCV1 -pend=5
    }
    #--- Collect motion:
    set fileIndexString $i
    set outputFile $logDir/$outputFileRoot$fileIndexString
    if $verbose {puts stdout "Collecting motion data..."}
    if [catch {CollectInjectedMotion -simulRun $simulRun -wmonFile $wmonFile -smonFile $smonFile \
		   -outputFile $outputFile -nTrajAver $average -plane $plane -stored $stored -verbose $verbose} result] {
	puts stdout "CollectInjectedMotion: $result"
    }
    set outputString "Motion rms: "
    foreach localPlane [list X Y] {
	if {[string first $localPlane $plane] != -1} {
	    if [catch {exec sddsprocess $outputFile.$localPlane.trajectory -pipe=out -filter=para,TurnAfterInj,0,0 \
			   -filter=col,s,$sStart,$sEnd -process=Trajectory,stand,Std \
			   | sdds2stream -pipe=in -para=Std} motionStDev$localPlane] {
		return -code error "Error processing $outputFile.$localPlane.trajectory: [set motionStDev$localPlane]"
	    }
	    append outputString "$localPlane: [set motionStDev$localPlane]  "
	}
    }
    puts stdout $outputString
}

if $doSetup {
    if [catch {LoadRAM -plane $plane -mode singleBunch -verbose $verbose} result] {
	puts stdout "LoadRAM: $result"
	exit
    }
}


puts stdout "Done."

