#!/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
}
source $locoBinDir/locoCommonProcedures

set usage "calculateAdjustTunes -lteFile <filename> -beamline <name> \[-paramFile <filename>\] -outputFile <filename> -Qx <number> -Qy <number> -knobFile <filename>  -tol <number> -outputStatusDevice <stdout|statusScreen|file>"

set args $argv
set lteFile ""
set beamline ""
set paramFile ""
set outputFile ""
set Qx ""
set Qy ""
set knobFile ""
set tol 0.001
set outputStatusDevice stdout
set outputStatusFile ""
APSParseArguments {lteFile beamline paramFile outputFile Qx Qy knobFile tol twiFile outputStatusDevice outputStatusFile}

if {![string length $lteFile] || ![string length $beamline] || ![string length $outputFile] || ![string length $Qx] \
	|| ![string length $Qy] || ![string length $knobFile]} {
    puts stderr $usage
    exit
}

#---------------------------------------------------------------------------------------------
# 01/23/25 11:49:50: needed to change the PV siffix from CurrentAO to
# PS:SetCurrentC
proc MakeParamFile {args} {
    APSParseArguments {filename knobFile xCoeff yCoeff}
    
    set tmpRoot /tmp/[APSTmpString]-tuneAdjust
    foreach pageName [list S:Tunex S:Tuney] coeff [list $xCoeff $yCoeff] \
	tmpFile [list $tmpRoot.x $tmpRoot.y] {
	    exec sddsconvert $knobFile -pipe=out \
		-rename=col,ControlName=ElementName -rename=col,Weight=ParameterValue \
		| sddsprocess -pipe=in $tmpFile -match=para,ControlName=$pageName \
		"-reedit=col,ElementName,%/:PS:SetCurrentC//" \
		"-redef=col,ParameterValue,ParameterValue $coeff *" \
		-print=col,ElementParameter,K1 -print=col,ParameterMode,DIFFERENTIAL \
		"-scan=col,QuadFamily,ElementName,%lf,edit=z: 2d" \
		"-redef=col,ParameterValue,QuadFamily 2 == pop pop ? ParameterValue : ParameterValue -1 * \$"
	}
    exec sddsxref $tmpRoot.x $tmpRoot.y -pipe=out -take=ParameterValue -match=ElementName \
	-rename=col,ParameterValue=ParameterValue1 \
	| sddsprocess -pipe "-redef=col,ParameterValue,ParameterValue ParameterValue1 +" \
	-redef=col,ElementOccurence,1,type=long \
	| sddsconvert -pipe=in $filename -del=col,ParameterValue1,QuadFamily
    file delete $tmpRoot.x $tmpRoot.y
}

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

proc GetTunes {args} {
    APSParseArguments {eleFile twiFile}
    if [catch {exec elegant $eleFile > $eleFile.log} result] {
        return -code error "Error running elegant: $result"
    }
    set Qx [exec sdds2stream $twiFile -para=nux]
    set Qy [exec sdds2stream $twiFile -para=nuy]
    return [list $Qx $Qy]
}

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

proc GetTuneDerivatives {args} {
    APSParseArguments {eleFile knobFile paramFile twiFile}

    OutputStatusMessage "Getting tune derivatives..."

    #--- List over non-perturbed tunes, x knob, and y knob:
    set coeff 0.01
    foreach coeffList [list "0 0" "$coeff 0" "0 $coeff"] listName [list tuneList0 tuneList1 tuneList2] {
        if [catch {MakeParamFile -filename $paramFile -knobFile $knobFile \
                     -xCoeff [lindex $coeffList 0] -yCoeff [lindex $coeffList 1]} result] {
            return -code error "MakeParamFile: $result"
        }
        if [catch {GetTunes -eleFile $eleFile -twiFile $twiFile} $listName] {
            return -code error "GetTunes: [set $listName]"
        }
    }
    set Qx1 [expr ([lindex $tuneList1 0] - [lindex $tuneList0 0]) / $coeff]
    set Qy1 [expr ([lindex $tuneList1 1] - [lindex $tuneList0 1]) / $coeff]
    set Qx2 [expr ([lindex $tuneList2 0] - [lindex $tuneList0 0]) / $coeff]
    set Qy2 [expr ([lindex $tuneList2 1] - [lindex $tuneList0 1]) / $coeff]
    return [list $Qx1 $Qy1 $Qx2 $Qy2]
}

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

proc AdjustAllSectors {args} {
    APSParseArguments {localParamFile knobFile eleTuneFile twiFile outputFile Qx Qy \
			   Qx1 Qx2 Qy1 Qy2 tol}

    if [catch {GetTunes -eleFile $eleTuneFile -twiFile $twiFile} tuneList] {
        return -code error "GetTunes: $tuneList"
    }
    set Qx0 [lindex $tuneList 0]
    set Qy0 [lindex $tuneList 1]

    set iter 5
    set done 0
    set I 0
    while {$done == 0} {
	set dQx [expr $Qx - $Qx0]
        set dQy [expr $Qy - $Qy0]
        set G1 [expr ($Qy2 * $dQx - $Qx2 * $dQy) / ($Qx1 * $Qy2 - $Qy1 * $Qx2)]
        set G2 [expr ($Qx1 * $dQy - $Qy1 * $dQx) / ($Qx1 * $Qy2 - $Qy1 * $Qx2)]
        if [catch {MakeParamFile -filename $localParamFile.iter -knobFile $knobFile \
		       -xCoeff $G1 -yCoeff $G2} result] {
            return -code error "MakeParamFile: $result"
        }
	exec sddsxref $localParamFile $localParamFile.iter -pipe=out -take=ParameterValue -match=ElementName \
	    -rename=col,ParameterValue=ParameterValue1 -fill -nowarning \
	    | sddsprocess -pipe "-redef=col,ParameterValue,ParameterValue ParameterValue1 +" \
	    | sddsconvert -pipe=in $localParamFile.tmp -del=col,ParameterValue1
        file copy -force $localParamFile.tmp $localParamFile
        if [catch {GetTunes -eleFile $eleTuneFile -twiFile $twiFile} tuneList] {
            return -code error "GetTunes: $tuneList"
        }
	set Qxprev $Qx0
	set Qyprev $Qy0
        set Qx0 [lindex $tuneList 0]
        set Qy0 [lindex $tuneList 1]
        if [catch {set dQx1 [expr abs($Qx0 - $Qxprev)]} result] {
	    OutputStatusMessage "Iteration $I, tune X: $result. Using nearest resonance tune."
	    set Qx0 [exec rpnl "$Qxprev 2 * 0.5 + int 2 /"]
	    set dQx1 [expr abs($Qx - $Qx0)]
	}
        if [catch {set dQy1 [expr abs($Qy0 - $Qyprev)]} result] {
	    OutputStatusMessage "Iteration $I, tune Y: $result. Using nearest resonance tune."
	    set Qy0 [exec rpnl "$Qyprev 2 * 0.5 + int 2 /"]
	    set dQy1 [expr abs($Qy - $Qy0)]
	}
	if {($dQx1 < $tol) && ($dQy1 < $tol)} {
	    set done 1
	} else {
	    incr I
	    if {$I > $iter} {
		OutputStatusMessage "Maximum number of iteration achieved."
		set done 1
	    }
	}
    }
    file copy -force $localParamFile $outputFile
    file delete $localParamFile.tmp $localParamFile.iter
    return $tuneList
}

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


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


#---------------------------------------------------------------------------------------------
proc OutputStatusMessage {text} {
    global outputStatusDevice outputStatusFile
    switch -exact $outputStatusDevice {
        statusScreen { APSSetVarAndUpdate status $text }
        stdout { puts stdout $text }
        file {
            if {![info exists outputStatusFile] || ![string length $outputStatusFile]} {
		#------ outputStatusFile is not defined. Send output to stdout...
		puts stdout $text
                return
            }
            if [file exists $outputStatusFile] {
                set fid [open $outputStatusFile a]
                puts $fid $text
                close $fid
            } else {
                set fid [open $outputStatusFile w]
                puts $fid $text
                close $fid
            }
        }
        default { return -code error "OutputStatusMessage: wrong outputStatusDevice: $outputStatusDevice" }
    }
}

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

    
set deleteFiles ""
set eleTuneFile apsTune.ele
set localParamFile apsTune.param
set twiFile apsTune.twi

lappend deleteFiles $eleTuneFile $localParamFile $eleTuneFile.log
if [string length $paramFile] {
    set parameterFileList [concat $paramFile $localParamFile]
} else {
    set parameterFileList $localParamFile
}
if [catch {Fit_GenerateElegantFileFromLTE1 -eleOutputFile $eleTuneFile \
	       -run_setup [list "lattice $lteFile use_beamline $beamline parameters $localParamFile.output"] \
	       -load_parameters [list "filename_list \"$parameterFileList\""] \
	       -twiss_output [list "filename $twiFile"] \
	       -run_control [list "n_steps 1"] \
	       -bunched_beam [list "n_particles_per_bunch 1"] \
	   } result] {
    puts stderr "Fit_GenerateElegantFileFromLTE: $result"
    exit
}
#------ Getting initial tunes...
if [catch {MakeParamFile -filename $localParamFile -knobFile $knobFile \
	       -xCoeff 0 -yCoeff 0} result] {
    puts stderr "MakeParamFile: $result"
    exit
}
if [catch {GetTunes -eleFile $eleTuneFile -twiFile $twiFile} tuneList] {
    puts stderr "GetTunes: $tuneList"
    exit
}
exec sddsprocess $localParamFile.output $localParamFile.output.0 -match=col,ElementType=*QUAD* \
    -match=col,ElementParameter=K1 "-match=col,ElementName=*:QS*,!"
lappend deleteFiles $localParamFile.output $localParamFile.output.0
set Qx0 [lindex $tuneList 0]
set Qy0 [lindex $tuneList 1]
OutputStatusMessage "Initial tunes are: Qx=[format %10.5f $Qx0], Qy=[format %10.5f $Qy0]"
#------ Getting tune derivatives...
if [catch {GetTuneDerivatives -eleFile $eleTuneFile -knobFile $knobFile \
	       -paramFile $localParamFile -twiFile $twiFile} derivativeList] {
    puts stderr "GetTuneDerivatives: $derivativeList"
    exit
}
set Qx1 [lindex $derivativeList 0]
set Qy1 [lindex $derivativeList 1]
set Qx2 [lindex $derivativeList 2]
set Qy2 [lindex $derivativeList 3]
#------ Return localParamFile to 0 after derivative calculation.
if [catch {MakeParamFile -filename $localParamFile -knobFile $knobFile \
	       -xCoeff 0 -yCoeff 0} result] {
    puts stderr "MakeParamFile: $result"
    exit
}

set Iterations 1
set done 0
set Iteration 0
while {$done == 0} {
    OutputStatusMessage "Iteration $Iteration."
    if [catch {AdjustAllSectors -localParamFile $localParamFile -knobFile $knobFile -eleTuneFile $eleTuneFile \
		   -twiFile $twiFile -outputFile $outputFile -Qx $Qx -Qy $Qy \
		   -Qx1 $Qx1 -Qx2 $Qx2 -Qy1 $Qy1 -Qy2 $Qy2 -tol [expr $tol / 10]} tuneList] {
	puts stderr "AdjustAllSectors: $tuneList"
	exit
    }

    set Qx0 [lindex $tuneList 0]
    set Qy0 [lindex $tuneList 1]
    if {[expr abs($Qx - $Qx0)] < $tol && [expr abs($Qy - $Qy0)]} {
	set done 1
    } else {
	file copy -force $outputFile $localParamFile
	OutputStatusMessage "Iteration $Iteration tunes are: Qx=[format %10.5f [lindex $tuneList 0]], Qy=[format %10.5f [lindex $tuneList 1]]"
	incr Iteration
    }
}
exec sddsprocess $outputFile -nowarning -redef=para,Qx0,$Qx0 -redef=para,Qy0,$Qy0 -redef=para,Qx,$Qx -redef=para,Qy,$Qy \
    -print=para,lteFile,$lteFile
lappend deleteFiles ${outputFile}~
eval file delete $deleteFiles
OutputStatusMessage "Done. Final tunes are: Qx=[format %10.5f [lindex $tuneList 0]], Qy=[format %10.5f [lindex $tuneList 1]]"

