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


source /home/helios/oagData/sr/commissioning/Scripts/commonProcedures

set usage "makeTuneAdjustParamFiles -lteFile <filename> -beamline <beamline-name> -targetTunes <nux nuy> -outputRoot <rootname> \[-familyQuadList <quad-list>\] \[-tol 0.0001\] \[-paramFileList <file-list>\] \[-removeKnobFiles 1\] \[-verbose 0\]"

set args $argv
set requiredArgList [list lteFile targetTunes outputRoot beamline]
set argList [concat $requiredArgList paramFileList verbose familyQuadList tol removeKnobFiles]
foreach arg $argList {set $arg ""}
set verbose 0
set quadFamilyList [list Q1 Q2]
set tol 0.0001
set removeKnobFiles 0
APSParseArguments $argList
set error 0
foreach arg $requiredArgList {if ![string length [set $arg]] {puts stdout "Missing argument: $arg"; set error 1}}
if $error {puts stderr $usage; exit 1}

set rootname tuneAdjustFiles
set eleFile $rootname.ele
set energy 6e3
set twiFile $rootname.twi
set quadDelta 0.001
set tuneDelta 0.01
set deleteFiles "$eleFile $rootname.param $rootname.twi"

#--------------------------------------------------------------------------------------------------------------
proc TuneError {args} {
    APSParseArguments {tuneList targetTunes}
    set dQ [expr sqrt(([lindex $tuneList 0] - [lindex $targetTunes 0])*([lindex $tuneList 0] - [lindex $targetTunes 0]) + \
			  ([lindex $tuneList 1] - [lindex $targetTunes 1])*([lindex $tuneList 1] - [lindex $targetTunes 1]))]
    return $dQ
}
#--------------------------------------------------------------------------------------------------------------

#------ Calculate tune response to quads:
set counter 0
foreach quadName [list S*:[lindex $quadFamilyList 0] S*:[lindex $quadFamilyList 0] S*:[lindex $quadFamilyList 1] ] k1Value [list 0 $quadDelta $quadDelta] {
    if [catch {Fit_GenerateElegantFileFromLTE1 -eleOutputFile $eleFile \
		   -run_setup [list "lattice $lteFile use_beamline $beamline default_order 2 p_central_mev $energy parameters $rootname.param"] \
		   -load_parameters [list "filename_list \"$paramFileList\""] \
		   -alter_elements [list "name $quadName item K1 value $k1Value differential 1"] \
		   -twiss_output [list "filename $twiFile"] \
		   -run_control [list "n_steps 1"] \
		   -bunched_beam [list "n_particles_per_bunch 1"] \
	       } result] {
	return -code error "Fit_GenerateElegantFileFromLTE1: $result"
    }
    catch {file delete $eleFile.log}
    if $verbose {puts stdout "elegant $eleFile > $eleFile.log"}
    if [catch {exec elegant $eleFile > $eleFile.log} result] {
	puts stdout "elegant returned error :$result"
	exit 1
    } else {file delete $eleFile.log}
    set tuneList [join [exec sdds2stream $twiFile -para=nux,nuy] ]
    if {$counter == 0} {
	set origTunes $tuneList
	puts stdout "Original tunes: [format %9.4f [lindex $origTunes 0]]; [format %9.4f [lindex $origTunes 1]]"
	exec sddsprocess $rootname.param $rootname.param.00 -match=col,ElementType=*QUAD* -match=col,ElementParameter=K1
	lappend deleteFiles $rootname.param.00
    } else {
	if {$counter == 1} {
	    set dx1 [expr ([lindex $tuneList 0] - [lindex $origTunes 0])/$quadDelta]
	    set dy1 [expr ([lindex $tuneList 1] - [lindex $origTunes 1])/$quadDelta]
	} else {
	    set dx2 [expr ([lindex $tuneList 0] - [lindex $origTunes 0])/$quadDelta]
	    set dy2 [expr ([lindex $tuneList 1] - [lindex $origTunes 1])/$quadDelta]
	}
    }
    incr counter
}

#------ Calculate quad knobs:
#------ Calculate quad changes:
exec sddsmakedataset $rootname.matrix -col=C1,type=double -data=$dx1,$dy1 -col=C2,type=double -data=$dx2,$dy2 
exec sddsmakedataset $rootname.vectorX -col=V1,type=double -data=$tuneDelta,0 
exec sddsmakedataset $rootname.vectorY -col=V1,type=double -data=0,$tuneDelta

exec sddsmatrixop $rootname.matrix $rootname.resX -inv -push=$rootname.vectorX -mult
exec sddsmatrixop $rootname.matrix $rootname.resY -inv -push=$rootname.vectorY -mult
set quadValueListx [exec sdds2stream $rootname.resX -col=doubleColumn0]
set quadValueListy [exec sdds2stream $rootname.resY -col=doubleColumn0]
lappend deleteFiles $rootname.matrix $rootname.vectorX $rootname.vectorY $rootname.resX $rootname.resY

#------ Read quad names:
set fileList ""
foreach quadName $quadFamilyList {
    exec sddsprocess $twiFile -pipe=out -match=col,ElementName=*$quadName -match=col,ElementType=*QUAD -print=col,TagName,%s#%1d,ElementName,ElementOccurence \
	| sddsconvert -pipe=in $rootname.quadNames.$quadName -retain=col,ElementName,TagName,ElementOccurence -del=para,*
    lappend fileList $rootname.quadNames.$quadName
}
set deleteFiles [concat $deleteFiles $fileList]

#------ Make ${plane}Up files:
foreach plane [list x y] {
    set fileList1 ""
    set quadValueList [set quadValueList$plane]
    foreach filename $fileList value $quadValueList {
	exec sddsprocess $filename $filename.tmp -print=col,ElementParameter,K1 -print=col,ParameterMode,DIFFERENTIAL \
	    "-redef=col,ParameterValue,$value"
	lappend fileList1 $filename.tmp
    }
    eval exec sddscombine $fileList1 -pipe=out -merge \
	| sddsprocess -pipe=in $outputRoot.${plane}Up0.01.param -redef=para,DeltaTune,$tuneDelta
    eval file delete $fileList1
}

#------ Make file for tune adjustment to target tunes:

if [catch {Fit_GenerateElegantFileFromLTE1 -eleOutputFile $eleFile \
	       -run_setup [list "lattice $lteFile use_beamline $beamline default_order 2 p_central_mev $energy parameters $rootname.param"] \
	       -load_parameters [list "filename_list \"$paramFileList\""] \
	       -load_parameters2 [list "filename_list \"<KNOBFILELIST>\""] \
	       -twiss_output [list "filename $twiFile"] \
	       -run_control [list "n_steps 1"] \
	       -bunched_beam [list "n_particles_per_bunch 1"] \
	   } result] {
    return -code error "Fit_GenerateElegantFileFromLTE1: $result"
}

set dQ [TuneError -tuneList $origTunes -targetTunes $targetTunes]

set tuneList $origTunes
set knobFileList ""
set iC 0
while {$dQ > $tol} {
    set dQx [expr [lindex $targetTunes 0] - [lindex $tuneList 0]]
    set dQy [expr [lindex $targetTunes 1] - [lindex $tuneList 1]]
    exec sddsprocess $outputRoot.xUp0.01.param $rootname.x.param.$iC "-redef=col,ParameterValue,ParameterValue $dQx $tuneDelta / *"
    exec sddsprocess $outputRoot.yUp0.01.param $rootname.y.param.$iC "-redef=col,ParameterValue,ParameterValue $dQy $tuneDelta / *"
    lappend knobFileList $rootname.x.param.$iC $rootname.y.param.$iC

    catch {file delete $eleFile.log}
    if $verbose {puts stdout "elegant $eleFile > $eleFile.log"}
    if [catch {exec elegant $eleFile "-macro=KNOBFILELIST=\"$knobFileList\"" > $eleFile.log} result] {
	puts stdout "elegant returned error :$result"
	exit 1
    } else {file delete $eleFile.log}
    set tuneList [join [exec sdds2stream $twiFile -para=nux,nuy] ]
    set dQ [TuneError -tuneList $tuneList -targetTunes $targetTunes]
    if $verbose {puts stdout "Tune error: [format %10.3e $dQ]"}
    incr iC
}
set finalTunes $tuneList
puts stdout "Final tunes:    [format %9.4f [lindex $finalTunes 0]]; [format %9.4f [lindex $finalTunes 1]]"



exec sddsprocess $rootname.param -pipe=out -match=col,ElementType=*QUAD* -match=col,ElementParameter=K1 \
    | sddsxref -pipe $rootname.param.00 -take=ParameterValue -rename=col,ParameterValue=ParameterValue0 \
    | sddsprocess -pipe "-redef=col,ParameterValue,ParameterValue ParameterValue0 -" -print=col,TagName,%s#%1d,ElementName,ElementOccurence \
    "-filter=col,ParameterValue,0,0,!" -print=col,ParameterMode,DIFFERENTIAL \
    -redef=para,Qx,[lindex $finalTunes 0] -redef=para,Qy,[lindex $finalTunes 1] \
    | sddsconvert -pipe=in $outputRoot.param -del=col,ParameterValue0,ParameterValueString,ElementGroup

file copy -force $rootname.twi $outputRoot.twi
eval file delete $deleteFiles $knobFileList

if $removeKnobFiles {
    file delete $outputRoot.xUp0.01.param $outputRoot.yUp0.01.param
    puts stdout "Files created: $outputRoot.param"
} else {
    puts stdout "Files created: $outputRoot.param, $outputRoot.xUp0.01.param, and $outputRoot.yUp0.01.param"
}
