#!/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)]

if [info exists env(LOCO_BINDIR)] {
    set locoBinDir $env(LOCO_BINDIR)
} else {
    puts stderr "Error: LOCO_BINDIR environment variable is not defined"
    exit
}
if [info exists env(LOCO_TMPDIR)] {
    set tmpDir $env(LOCO_TMPDIR)
} else {
    set tmpDir $env(HOME)/tmp
}

source $locoBinDir/locoCommonProcedures
source $locoBinDir/locoFitProcedures

set requiredVariableList [list latticeInputFile beamlineName hrmMeasuredInput vrmMeasuredInput \
			      elementListVarFile elementListRMFile optionsList \
			      workDir abortFile corFraction iterations \
			      resultsFile definitionFile]

foreach var $requiredVariableList {set $var ""}
set args $argv
APSParseArguments $requiredVariableList
set missingVarList ""
foreach var $requiredVariableList {
    if ![string length [set $var]] {lappend missingVarList $var}
}
if [llength $missingVarList] {
    puts stderr "locoFitting: The following variables are missing: $missingVarList"; exit
}

#------ Not required variables:
set nonrequiredVariableListString [list paramFileList outputStatusDevice outputStatusFile \
				       previousResultsFile acceleratorCode calculateEmittanceList \
				       machineName kickFileX kickFileY rawRMDeriv]
set nonrequiredVariableListNumber [list measuredFormat usePreviousSolution useBpmWeight analMode \
				       averLevel nonlinLevel badPointsLevel SVnumber nux nuy \
				       recalculateDerivative filterDerivative recalculateInverse recombineDerivative prepareMatricesOnly \
				       useKickFiles verbose calculateDerivOnly bigIterNumber]

foreach var $nonrequiredVariableListString {set $var ""}
set outputStatusDevice stdout
set machineName APS
set acceleratorCode elegant
foreach var $nonrequiredVariableListNumber {set $var 0}
set bigIterNumber 3
set analMode 2
APSParseArguments [concat $nonrequiredVariableListString $nonrequiredVariableListNumber]
set checkEmptyVars ""
if $usePreviousSolution {lappend checkEmptyVars previousResultsFile}
if $useKickFiles {lappend checkEmptyVars kickFileX kickFileY}
set missingVarList ""
foreach var $checkEmptyVars {
    if ![string length [set $var]] {lappend missingVarList $var}
}
if [llength $missingVarList] {
    puts stderr "locoFitting: The following variables are missing: $missingVarList"; exit
}

array set optionsArray $optionsList
if $optionsArray(useTune) {
    if {$nux ==0 || $nuy == 0} {puts stderr "Error: No tunes given with tune fitting."; exit}
}

#------ Make some files local (it is possible that individual nodes don't have access to remote file storage):
set copyfileList [list latticeInputFile]
if $usePreviousSolution {lappend copyfileList previousResultsFile}
foreach filenameVar $copyfileList {
    if {[string compare [set $filenameVar] $workDir/[file tail [set $filenameVar]]] != 0} {
	#--- file copy hard to figure out with links, cp is easier
	#--- cp gives error if the target a broken link; "file exists" gives 0 for broken link, so just do "catch {delete}"
	catch {file delete $workDir/[file tail [set $filenameVar]]}
	exec cp -f [set $filenameVar] $workDir/[file tail [set $filenameVar]]
	set $filenameVar $workDir/[file tail [set $filenameVar]]
    }
}
if [string length $paramFileList] {
    set newParamFileList ""
    foreach filename $paramFileList {
	if {[string compare $filename $workDir/[file tail $filename]] != 0} {
	    file copy -force $filename $workDir/[file tail $filename]
	    lappend newParamFileList $workDir/[file tail $filename]
	} else {
	    lappend newParamFileList $filename
	}
    }
    set paramFileList $newParamFileList
}

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

proc ReadListFromFile { args } {
    set optionsList ""
    APSParseArguments { listFile optionsList }
    set mainList ""
    if [string length $optionsList] {
	array set options $optionsList
	set indexList [list useQuads useCorrs useCorrs useBPMs useBPMs useEnergy useSkew useCorrsTilt useCorrsTilt \
	    useBPMsTilt useBPMsTilt]
	set paramList [list Quad HCorr VCorr HBPM VBPM Energy Skew HCorrTilt VCorrTilt HBPMTilt VBPMTilt]
	foreach index $indexList param $paramList {
	    if $options($index) {
		set localList [join [exec sddsprocess $listFile -pipe=out -match=para,ListName=$param -nowarning \
					 | sdds2stream -pipe=in -col=List] ]
		lappend mainList $localList
	    } else {
		lappend mainList ""
	    }
	}
    } else {
	foreach param [list HCorr VCorr HBPM VBPM] {
	    lappend mainList [join [exec sddsprocess $listFile -pipe=out -match=para,ListName=$param -nowarning \
					| sdds2stream -pipe=in -col=List] ]
	}
    }
    return $mainList
}

#-----------------------------------------------------------------------------------------------------------------------
proc GetBadBPMs {args} {
    APSParseArguments {useBpmWeight}
    global workDir tmpDir definitionFile
    if [catch {Fit_ReadDefinitionFile -definitionFile $definitionFile -workDir $workDir -warnings 0 \
	       -variableList "dispColumn hrmDifference vrmDifference weightFile badBpmSigmaLimit"} result] {
	return -code error "Fit_ReadDefinitionFile: $result"
    }
    set tmpRoot $tmpDir/[APSTmpString]-badBpms

    #--- Bad bpms:
    if $useBpmWeight {
	if [catch {exec sddsxref $hrmDifference $vrmDifference -pipe=out -take=* -leave=TagName \
		       | sddsxref -pipe $weightFile -take=Weight -rename=col,Weight=WeightName \
		       | sddsprocess -pipe "-redef=col,%s,%s WeightName *,select=*,exclude=*Name" \
		       | sddsconvert -pipe -del=col,$dispColumn,WeightName \
		       | sddsrowstats -pipe -standard=StdDev,* \
		       | sddsoutlier -pipe=in $tmpRoot -col=StdDev -stDevLimit=$badBpmSigmaLimit -invert -nowarning} result] {
	    return -code error "Error finding outliers: $result"
	}
    } else {
	if [catch {exec sddsxref $hrmDifference $vrmDifference -pipe=out -take=* -leave=TagName \
		       | sddsconvert -pipe -del=col,$dispColumn \
		       | sddsrowstats -pipe -standard=StdDev,* \
		       | sddsoutlier -pipe=in $tmpRoot -col=StdDev -stDevLimit=$badBpmSigmaLimit -invert -nowarning} result] {
	    return -code error "Error finding outliers: $result"
	}
    }
    set badBpmsH [join [exec sddsconvert $tmpRoot -pipe=out -keep=1 \
			    | sdds2stream -pipe=in -col=TagName] ]
    set badBpmsV [join [exec sddsconvert $tmpRoot -pipe=out -keep=2 \
			    | sdds2stream -pipe=in -col=TagName] ]
    puts stdout "The following BPMs will be set to low weight (badBpmSigmaLimit: $badBpmSigmaLimit): H plane: $badBpmsH; V plane: $badBpmsV"

    #--- Bad correctors:
    if [catch {exec sddsxref $hrmDifference $vrmDifference -pipe=out -take=* -leave=TagName \
		   | sddsconvert -pipe -del=col,$dispColumn \
		   | sddscombine -pipe -merge \
		   | sddstranspose -pipe -oldcolumn=Corrector \
		   | sddsrowstats -pipe -standard=StdDev,* \
		   | sddsoutlier -pipe -col=StdDev -stDevLimit=$badBpmSigmaLimit -invert -nowarning \
		   | sdds2stream -pipe=in -col=Corrector} badCorrectors] {
	return -code error "Error finding outliers: $badCorrectors"
    }
    if [llength $badCorrectors] {puts stdout "Warning: the following correctors are bad: [join $badCorrectors ]."}
    file delete $tmpRoot
    return [list H $badBpmsH V $badBpmsV]
}
#-----------------------------------------------------------------------------------------------------------------------
proc MakeBpmWeightFile {args} {
    global tmpDir workDir definitionFile
    APSParseArguments {badBpmList useBpmWeight}
    array set badBpmArray $badBpmList
    set useFile 0
    set bpmList ""
    set columnList ""
    if [llength $badBpmArray(H)] {
	set useFile 1
	foreach bpm $badBpmArray(H) {
	    lappend bpmList $bpm
	    lappend planeList X
	}
    }
    if [llength $badBpmArray(V)] {
	set useFile 1
	foreach bpm $badBpmArray(V) {
	    lappend bpmList $bpm
	    lappend planeList Y
	}
    }
    if $useFile {
	set tmpRoot $tmpDir/[APSTmpString]
	if [catch {Fit_ReadDefinitionFile -definitionFile $definitionFile -workDir $workDir -warnings 0 \
		       -variableList "manualWeightFile"} result] {
	    return -code error "Fit_ReadDefinitionFile: $result"
	}
	exec sddsmakedataset -pipe=out \
	    -col=TagName,type=string -data=[join $bpmList ,] \
	    -col=Plane,type=string -data=[join $planeList ,] \
	    | sddsprocess -pipe=in $tmpRoot -redef=col,Weight,0.01
	if $useBpmWeight {
	    exec sddscombine $manualWeightFile $tmpRoot -pipe=out -merge -overWrite \
		| sddsconvert -pipe=in $tmpRoot.1 -ascii
	    file copy -force $tmpRoot.1 $manualWeightFile
	    file delete $tmpRoot $tmpRoot.1
	} else {
	    exec sddsconvert $tmpRoot $manualWeightFile -ascii
	    file delete $tmpRoot
	}
    }
    return $useFile
}
#-----------------------------------------------------------------------------------------------------------------------

array set optionsArray $optionsList
###
set optionsArray(fitPathLength) 0
###

if [catch {Fit_ReadListFromFile -listFile $elementListRMFile} elementListRM] {
    puts stderr "Fit_ReadListFromFile (1): $elementListRM"
    exit
}
if [catch {Fit_ReadListFromFile -listFile $elementListVarFile -optionsList [array get optionsArray]} elementListVar] {
    puts stderr "Fit_ReadListFromFile (2): $elementListVar"
    exit
}
if {$iterations == 0} {set bigIterNumber 1}
#------ badPointsLevel is unset inside CalculateLOCOFit, so we remember it here.
set badPointsLevel1 $badPointsLevel
for {set II 0} {$II < $bigIterNumber} {incr II} {
    if [catch {CalculateLOCOFit \
		   -bigIterNumber $bigIterNumber \
		   -bigIter $II \
		   -acceleratorCode $acceleratorCode \
		   -latticeInputFile $latticeInputFile \
		   -latticeParamFileList $paramFileList \
		   -beamlineName $beamlineName \
		   -hrmMeasuredInput $hrmMeasuredInput \
		   -vrmMeasuredInput $vrmMeasuredInput \
		   -rawRMDeriv $rawRMDeriv \
		   -elementListVar $elementListVar \
		   -elementListRM $elementListRM \
		   -optionsList [array get optionsArray] \
		   -nuxMeasured $nux \
		   -nuyMeasured $nuy \
		   -workDir $workDir \
		   -calculateDerivOnly $calculateDerivOnly \
		   -usePreviousSolution $usePreviousSolution \
		   -previousResultsFile $previousResultsFile \
		   -recalculateDerivative $recalculateDerivative \
		   -filterDerivative $filterDerivative \
		   -recombineDerivative $recombineDerivative \
		   -recalculateInverse $recalculateInverse \
		   -useMeasuredDisp 0 \
		   -continuePrevious 0 \
		   -useBpmWeight $useBpmWeight \
		   -SVnumber $SVnumber \
		   -corFraction $corFraction \
		   -iterations $iterations \
		   -averLevel $averLevel \
		   -nonlinLevel $nonlinLevel \
		   -badPointsLevel $badPointsLevel \
		   -resultsFile $resultsFile \
		   -definitionFile $definitionFile \
		   -coupledMatrix $analMode \
		   -useKickFiles $useKickFiles \
		   -kickFileX $kickFileX \
		   -kickFileY $kickFileY \
		   -prepareMatricesOnly $prepareMatricesOnly \
		   -outputStatusDevice $outputStatusDevice \
		   -verbose $verbose \
		   -abortFile $abortFile
    } resErrorList] {
	puts stderr "CalculateLOCOFit: $resErrorList"
	exit
    }

    if {$prepareMatricesOnly || $calculateDerivOnly} {break}

    #------ Element lists are changed inside CalculateLOCOFit (Dispersion added), so re-read here:
    if [catch {Fit_ReadListFromFile -listFile $elementListRMFile} elementListRM] {
	puts stderr "Fit_ReadListFromFile (1): $elementListRM"
	exit
    }
    if [catch {Fit_ReadListFromFile -listFile $elementListVarFile -optionsList [array get optionsArray]} elementListVar] {
	puts stderr "Fit_ReadListFromFile (2): $elementListVar"
	exit
    }

    #------ Generate bad bpms if not last bigIter:
    set useWeightFile 0
    if {$II != [expr $bigIterNumber - 1]} {
	if {$badPointsLevel1 != 0} {
	    if [catch {GetBadBPMs -useBpmWeight $useBpmWeight} badBpmList] {
		puts stderr "GetBadBPMs: $badBpmList"
		exit
	    }
	    if [catch {MakeBpmWeightFile -badBpmList $badBpmList -useBpmWeight $useBpmWeight} useWeightFile] {
		puts stderr "MakeBpmWeightFile: $useWeightFile"
		exit
	    }
	}
    }
    file copy -force $resultsFile $resultsFile.$II
    file copy -force $iterationsLogFile $iterationsLogFile.$II

    if $useWeightFile {set useBpmWeight 1}
    set usePreviousSolution 1
    set previousResultsFile $resultsFile.$II
    set badPointsLevel $badPointsLevel1
    set recalculateDerivative 0
    set filterDerivative 0
    set recombineDerivative 0
    set recalculateInverse 0
}

#------ Calculate emittances:
if {[string length $calculateEmittanceList] && $iterations != 0 && ![file exists $abortFile]} {
    puts stdout "\nCalculating emittances using beam envelopes..."
    if [llength $paramFileList] {set paramFileListString " -paramFileList \"$paramFileList\""} else {set paramFileListString ""}
    array set emitArrayList $calculateEmittanceList
    if [catch {eval CalculateEmittances -workDir $workDir -definitionFile $definitionFile -lteFile $latticeInputFile $paramFileListString \
		   -resultsFile $resultsFile -beamline $emitArrayList(beamline) -rfVoltageMV $emitArrayList(rfVoltageMV) -rfH $emitArrayList(rfH) \
		   -beamEnergyGeV $emitArrayList(beamEnergyGeV)} result ] {
	puts stdout "CalculateEmittances: $result"
    }
}
exit

