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

#--- LOCO_BINDIR_OUTSIDE is used to redefine LOCO_BINDIR inside qsub command. Somehow direct redefinition of LOCO_BINDIR in qsub command does not work
if [info exists env(LOCO_BINDIR_OUTSIDE)] {
    set env(LOCO_BINDIR) $env(LOCO_BINDIR_OUTSIDE)
}

if [info exists env(LOCO_BINDIR)] {
    set locoBinDir $env(LOCO_BINDIR)
} else {
    puts stderr "Error: LOCO_BINDIR environment variable is not defined"; exit
}
source $locoBinDir/locoCommonProcedures

set args $argv
set usage "usage: quadsOnly -argMode <file|direct> \[-scriptParametersFile <filename>\] \[-verbose <0|1>\] \[-deleteFilesOnExit 0\]"

set argMode ""
set verbose 0
set deleteFilesOnExit 1
APSParseArguments {argMode verbose deleteFilesOnExit}
if ![string length $argMode] {puts stdout $usage; exit}
set verbose1 $verbose

if {[string compare $argMode "file"] == 0} {
    set scriptParametersFile ""
    APSParseArguments {scriptParametersFile}
    if {![string length $scriptParametersFile]} {puts stderr $usage; exit}
    if [catch {open $scriptParametersFile r} fid] {return -code error "$fid"} else {gets $fid args; close $fid}
}

set argList [list path lteFile beamlineName latticeParamFileList optimFile varList tuneInclude parameterDelta \
		 dispColumn dispFitWeight useDoubleOrbits verbose \
		 kick fixedOrbitLength directRMCalc fitDispersion computedRMX computedRMY optimNonlin\
		 coupledMatrix bRho makeBaselineCalc locoSearchString closedOrbit \
		 twissRefFile twissRefElement malignElement localTmpDir useKickFiles kickFileX kickFileY bpmMode \
		 specialElementsInputFile variableFile]
foreach arg $argList {set $arg ""}
set makeBaselineCalc 0
set bpmMode 0
APSParseArguments $argList
if {[string length $verbose] == 0} {set verbose $verbose1}

if $verbose {puts stdout "locoRunQuads: LOCO_BINDIR: $locoBinDir"}

#-------------------------------------------------------------------------------------------------------------------------
proc MakeElementNameFromTagName_Old {args} {
    set tagNameColumn TagName
    set output ""
    APSParseArguments {input output tagNameColumn}
    if {[string length $output] == 0} {set localOutput $input.tmp} else {set localOutput $output}
    #------ elegant can create a mixture of names when some names have #N in it and others don't.
    #------ To handle it, add #1 at the end, then delete second # symbol and after it.
    exec sddsprocess $input -pipe=out "-edit=col,ElementName,$tagNameColumn,S?@#@ K" \
	"-edit=col,TagNameCopy,$tagNameColumn,e i@#1@ a 2S?@#@ 2d" \
	"-scan=col,ElementOccurence,TagNameCopy,%ld,type=long,edit=Z#" \
	| sddsconvert -pipe=in $localOutput -del=col,TagNameCopy
    if {[string length $output] == 0} {file copy -force $localOutput $input; file delete $localOutput}
}
#-------------------------------------------------------------------------------------------------------------------------

## ----> Initial preparations before the Loop ---------------------------------------------------------------------

set hostname [info hostname]
if [catch {Fit_WriteToFile -filename $path/hostname -accessMode w -line $hostname} result] {
    return -code error $result
}

#------ Creating local (on a node) tmp dir (including all upper directories):
#------ Here path is .../quads/NN type and localTmpDir is .../quads. So we copy last subdir to avoid file conflicts.
set tmpDir $localTmpDir/[file tail $path]
file mkdir $tmpDir

set latticeRMFile $path/latticeRM.file

if [catch {Fit_DeleteElementsFromList -elementList [join [exec sddsquery -col $computedRMX] ] \
               -deleteElements "BPMName TagName $dispColumn"} hCorrList] {
    return -code error "Fit_DeleteElementsFromList: $hCorrList"
}
if [catch {Fit_DeleteElementsFromList -elementList [join [exec sddsquery -col $computedRMY] ] \
               -deleteElements "BPMName TagName $dispColumn"} vCorrList] {
    return -code error "Fit_DeleteElementsFromList: $vCorrList"
}
set hYes [llength $hCorrList]
set vYes [llength $vCorrList]

if {$tuneInclude} {
    set nux0 [exec sdds2stream -para=nux $computedRMX]
    set nuy0 [exec sdds2stream -para=nuy $computedRMY]
}

set delta $parameterDelta

#--- makeBaselineCalc means that computedRM[XY] files are ignored and will be calculated at the begining.
if $makeBaselineCalc {
    set varList [linsert $varList 0 [lindex $varList 0]]
}

## ----> The main Loop. ------------------------------------------------------------------------------------------

if $verbose { puts stdout "[clock format [clock seconds] -format %H:%M:%S]: Loop is beginning: $varList" }

if !$bpmMode {set localParamFileList [concat $latticeParamFileList $tmpDir/local.param]} else {set localParamFileList $latticeParamFileList}
set tmpRootMain $tmpDir/[APSTmpString]
foreach elementName $varList {
    set deleteFiles ""
    if $verbose { puts stdout "[clock format [clock seconds] -format %H:%M:%S]: Working on $elementName" }
    set tmpRoot $tmpRootMain.$elementName
    set RMdifferenceFile  $tmpRoot.dif
    set RMXdifferenceFile $tmpRoot.xdif
    set RMYdifferenceFile $tmpRoot.ydif
    set sddstmpFileX $tmpRoot.2
    set sddstmpFileY $tmpRoot.3
    set hrmFile $tmpRoot.hResponse.hrm
    set vrmFile $tmpRoot.vResponse.vrm
    set twiFile $tmpRoot.twi
    lappend deleteFiles $RMdifferenceFile $RMXdifferenceFile $RMYdifferenceFile $sddstmpFileX $sddstmpFileY \
	$hrmFile $vrmFile $hrmFile.XRM $vrmFile.XRM $twiFile
    set columnName $elementName
    if $makeBaselineCalc {
        set delta1 0
    } else {
        set delta1 $delta
    }
    #--- Creating param file with changed variable to run calculateResponseMatrix:
    if !$bpmMode {
	if {[string length $specialElementsInputFile] == 0} {
	    if [catch {exec sddsprocess $variableFile $tmpDir/local.param -match=col,TagName=$elementName -redef=col,ParameterValue,$delta1 \
			   "-scan=col,ElementOccurence,TagName,%ld,type=long,edit=Z#"} result] {
		return -code error "Error making local.param file from $variableFile (element $elementName probably missing): $result"
	    }
	} else {
	    if !$makeBaselineCalc {set delta1 [exec sddsprocess $specialElementsInputFile -pipe=out -match=col,TagName=$elementName | sdds2stream -pipe=in -col=Delta]}
	    set suffix [exec sddsprocess $specialElementsInputFile -pipe=out -match=col,TagName=$elementName | sdds2stream -pipe=in -col=NameSuffix]
	    if [catch {HandleSpecialVariableType -specialElementsInputFile $specialElementsInputFile -specialTagName $elementName -specialDelta $delta1 -outputFile $tmpDir/local.param} result] {
		return -code error "HandleSpecialVariableType: $result"
	    }
	}
	set bpmString ""
    } else {
	foreach ext [list xBpm yBpm bpmTilt bpmCoef] {
	    exec sddsprocess [file root $variableFile].$ext $tmpRoot.$ext -redef=col,ParameterValue,0 -nowarning
	}
	exec sddsprocess $variableFile $tmpRoot[file extension $variableFile] \
	    "-redef=col,ParameterValue,TagName \"$elementName\" streq ? $delta1 : 0 \$"
	set bpmString " -bpmGainXFile $tmpRoot.xBpm -bpmGainYFile $tmpRoot.yBpm -bpmTiltFile $tmpRoot.bpmTilt -bpmCoefFile $tmpRoot.bpmCoef "
	lappend deleteFiles $tmpRoot.xBpm $tmpRoot.yBpm $tmpRoot.bpmTilt $tmpRoot.bpmCoef
    }

    if $useKickFiles {set kickFileString " -useKickFiles $useKickFiles -kickFileX $kickFileX -kickFileY $kickFileY "} else {set kickFileString ""}
    if !$closedOrbit {set twissRefString " -twissRefFile $twissRefFile -twissRefElement $twissRefElement "} else {set twissRefString ""}
    set rmCommand "$locoBinDir/locoCalculateResponseMatrix \
                   -runMode getResponseMatrix \
                   -lteFile $lteFile \
                   -beamlineName $beamlineName \
                   -latticeParamFileList \"$localParamFileList\" \
                   -optimFile $latticeRMFile \
                   -hrmFile $hrmFile \
                   -vrmFile $vrmFile \
                   -hCorrList \"$hCorrList\" \
                   -vCorrList \"$vCorrList\" \
                   -workDir $tmpDir \
                   -useQsub 0 \
                   -bRho $bRho \
                   -twiFile $twiFile \
                   -dispFitWeight $dispFitWeight\
                   -directCalc $directRMCalc \
                   -kickX $kick \
                   -kickY $kick \
                   $kickFileString \
                   -fixedOrbitLength $fixedOrbitLength \
                   -fitDispersion $fitDispersion \
                   -dispColumn $dispColumn \
                   -coupledMatrix $coupledMatrix \
                   -closedOrbit $closedOrbit \
		   $twissRefString \
		   -malignElement $malignElement \
                   -useDoubleOrbits $useDoubleOrbits \
                   $bpmString \
                   -verbose $verbose"
    if $verbose {
	puts stdout "[clock format [clock seconds] -format %H:%M:%S]: Running locoCalculateResponseMatrix:"
	puts stdout $rmCommand
    }
    if [catch {eval exec $rmCommand >@ stdout} result] {
        return -code error "locoCalculateResponseMatrix: $result"
    }
    if $verbose {puts stdout "[clock format [clock seconds] -format %H:%M:%S]: locoCalculateResponseMatrix done."}

    #--- Leave only #X in page 1 and #Y in page 2:
    foreach rmFile [list $hrmFile $vrmFile] {
	exec sddsconvert $rmFile -pipe=out -keep=1 \
	    | sddsprocess -pipe=in $tmpRoot.rm1 -match=col,TagName=*#X*
	exec sddsconvert $rmFile -pipe=out -keep=2 \
	    | sddsprocess -pipe=in $tmpRoot.rm2 -match=col,TagName=*#Y*
	exec sddscombine $tmpRoot.rm1 $tmpRoot.rm2 $rmFile -overwrite
	file delete $tmpRoot.rm1 $tmpRoot.rm2
    }

    #------ Dispersion is not changed by CALIBRATION in parameter file, so we need to change it manually here:
    if {[string compare $dispColumn $elementName] == 0} {
	exec sddsprocess $hrmFile $tmpRoot.dispCal "-redef=col,$dispColumn,$dispColumn 1 $delta1 + *"
	file copy -force $tmpRoot.dispCal $hrmFile
	file delete $tmpRoot.dispCal
    }

    if $makeBaselineCalc {
        set planeFlagList [list $hYes $vYes]
        set rmFileList [list $hrmFile $vrmFile]
        set compFileList [list computedRMX computedRMY]
        set extList [list RMX RMY]
        foreach planeFlag $planeFlagList rmFile $rmFileList compFile $compFileList ext $extList {
            if $planeFlag {
                exec sddscombine $rmFile -pipe=out -merge \
                    | sddsselect -pipe [set $compFile] -match=TagName \
                    | sddsconvert -pipe=in $tmpRoot.$ext -del=col,s
                set $compFile $tmpRoot.$ext
            }
        }
        if {$tuneInclude} {
            set nux0 [exec sdds2stream -para=nux $computedRMX]
            set nuy0 [exec sdds2stream -para=nuy $computedRMY]
        }
        set makeBaselineCalc 0
    } else {
        set planeFlagList [list $hYes $vYes]
        set rmFileList [list $hrmFile $vrmFile]
        set compFileList [list $computedRMX $computedRMY]
        set diffFileList [list $RMXdifferenceFile $RMYdifferenceFile]
        set colFileList [list $tmpRoot.colx $tmpRoot.coly]
	lappend deleteFiles $tmpRoot.colx $tmpRoot.coly
        set nuList [list nux nuy]
        foreach planeFlag $planeFlagList rmFile $rmFileList compFile $compFileList diffFile $diffFileList \
            columnFile $colFileList nu $nuList {
                if $planeFlag {
                    exec sddscombine $rmFile -pipe=out -merge \
                        | sddsselect -pipe $compFile -match=TagName \
                        | sddsconvert -pipe -del=col,s \
                        | sddschanges -pipe -base=$compFile -changes=exclude=*Name,* -copy=BPMName,TagName \
                        | sddsconvert -pipe "-editNames=col,Change*,a 8d" \
                        | sddsprocess -pipe=in $diffFile "-redefine=col,%s,%s $delta1 /,select=*,exclude=*Name,units=m/rad"
                    if [catch {Fit_MakeColumnFromMatrix -matrixFile $diffFile -columnFile $columnFile \
                                   -columnName $columnName -rowNameColumn TagName -deleteColumns BPMName} result] {
                        return -code error "Fit_MakeColumnFromMatrix: $result"
                    }
                } else {
                    exec sddsmakedataset $columnFile -col=Rootname,type=string -data -col=$columnName,type=double -data
                }
            }
        set singleColumnFile $tmpRoot.$elementName
        eval exec sddscombine $colFileList $singleColumnFile -overWrite -merge
        if $tuneInclude {
            #--- lindex in case the file is 2-page.
	    if {[lsearch -exact [exec sdds2stream -para=nux $hrmFile] NaN] != -1} {
		return -code error  "Error in tune calculations. Check that your parameterDelta is correct: $parameterDelta"
	    }
            set nuxDif [expr ([lindex [exec sdds2stream -para=nux $hrmFile] 0] - $nux0) / $delta1]
            set nuyDif [expr ([lindex [exec sdds2stream -para=nuy $vrmFile] 0] - $nuy0) / $delta1]
            exec sddsmakedataset $tmpRoot.tunes -col=Rootname,type=string -data=Snux,Snuy \
                -col=$elementName,type=double -data=$nuxDif,$nuyDif
            exec sddscombine $singleColumnFile $tmpRoot.tunes $tmpRoot.1 -overWrite -merge
            file copy -force $tmpRoot.1 $singleColumnFile
            lappend deleteFiles $tmpRoot.tunes $tmpRoot.1
        }
	if [string length $specialElementsInputFile] {
	    exec sddsconvert $singleColumnFile -nowarning -rename=col,$elementName=$elementName$suffix
	    file delete ${singleColumnFile}~
	}
        lappend RMDFileList $singleColumnFile
    }
    eval file delete $deleteFiles
}

if $verbose { puts stdout "Loop is done..." }

#------ Combining all columns after the loop...

if {[llength $RMDFileList] == 1} {
    file copy -force $RMDFileList $path/responseMatrix.sdds
} else {
    eval exec sddsxref $RMDFileList $path/responseMatrix.sdds -take=* -leave=Rootname -nowarning
}
eval file delete $computedRMX $computedRMY $tmpDir/local.param $RMDFileList

if $verbose {puts stdout "Matrix is assembled..."}

if $deleteFilesOnExit {catch {file delete -force $tmpDir}}

if [catch {open $path/task.completed w} fid] {
    puts stderr "Error $fid"
    return -code error $fid
} else {
    puts $fid "END"
    close $fid
}

exit
