; Copyright (c) 1998-2010 A.P. Hitchcock  All rights reserved
;+
;NAME:
;	  STACK_FIT
;
;LAST CHANGED: ----------------------------------- 20-Jul-10
;
;PURPOSE:
;	This procedure selects inputs for fitting the spectrum at each pixel
; in a stack to sum of model spectra
; It uses COMMON blocks to store and transfer data
;
;CATEGORY:
;	AXIS:  stack analysis (operates stand-alone as well)
;
;CALLING SEQUENCE:
; for procedures:
;	STACK_FIT(/AXIS,/verbose)
;
;CALLED FROM AXIS:
;	->Stacks->stack_fit
;
;INPUTS:
;	All input parameters are obtained through user dialogs.
;
;KEYWORDS:
;	AXIS - if set, indicates called from AXIS
;	VERBOSE - if set, displays results as generated line-by-line (SLOW!)
;
;ROUTINES
;	STACK_FIT - main routine to execute stack fit
;	GROOM1 - linear fit to pre-edge values
;	INPUTS to GROOM1
;	spectrum - the input and the result after subtraction of back
;	back     - the constant to be subtracted
;	integral - integrated value in background region
;	zero     - number of points to evaluate constant background
;
;OUTPUTS:
;	If run from AXIS, the results are placed in AXIS buffers
; buffer 1 - n are the spectra of the n-models to be used in the fit,
;              interpolated to the energy values of the stack.
; buffer 4 - (4+n) are the resulting composition maps for the n components
; buffer 8 - linear term of fit
; buffer 9 - chi-square of the fit
; otherwise display windows are generated as required
;
;COMMON BLOCKS:
;	@AXIS_COM	standard set of common blocks
;	@stack_process_com
;	Volume_data - STACK
; 	@BSIF_COM - nc data
;
;PROCEDURE:
;	After reading in the stack (*.nc) and the model spectra (*.txt)
; and the names to be used for the output files,
; the spectrum at each pixel is fit to (linear) + SUM{coeff(i)*model spectra(i)}
; The coefficients at each pixel constitute the component map for species i.
; The user can choose to perform a pre-edge subtraction based on the first {m} -data points
; in order to base the analysis on solely the core spectra signal
;
;MODIFICATION HISTORY:
; ***********
; ORIGINAL kernel of this routine was written by Rick Kneedler
; ***********
; (25-May-99 aph) first developed from Kneedler code
; (19-Jun-99 aph) allow user choice to background subtract image spectra
; (07-jul-99 aph) extend to arbitrary number of components
; (09-sep-99 aph) store all results automatically; correct (x,y) scale stuff
; (20-sep-99 aph) remove linear term; set NaN, Inf to zero; integer col #
; (23-oct-99 aph) force plot of components on store
; ( 7-oct-99 aph) allow assignment of simple names to output files
; (26-nov-99 aph) get rid of error message for text prompt
; (23-Dec-99 aph) correct bad point filter on output component maps
; (27-feb-00 aph) add groupIP to get_text cal; axis standard documentation
; (09-apr-00 aph) change groom.pro to fit a sloped line with feedback
; (19-oct-00 aph) increased max # of components to 8
; (13-nov-00 aph) fixed error for GE 5 component; model data
; (09-jan-01 aph) move groom to run with compile-on-fly
; (12-aug-01 aph) put overwrite checking into file writing procedure
; (14-may-03 aph) use comment line, not filename for prompt for name
;					 par filesaves filenames and labels only
; (28-may-03 aph) fix-up no par; force extension to 'par'; fix up comp_names
; (04-jun-03 aph) use ax_par_save, ax_par_load
; (11-jun-03 aph) fix 1 boo-boo
; (30-dec-03 aph) add ax_sort_mono to force stack E-scale and reference spectral scales to be monotonic
; (20-jan-04 aph) use dialog_message to control residual stack writing
;				  ensure E-value (0.0) added to component maps
; (10-Jun-04 aph) remove pre-edge background subtraction
; (01-feb-05 aph) fix ax_sort_mono to use structure (ricochet change from lox change)
; (17-Sep-08 aph) replace analcom with stack_process_com
; (30-jul-09 aph) add stack file name to component map names
;                 truncate short name for component maps at the first space
; (20-Jul-10 aph) replace nwin,0 with window, 0  (line 247) - code error                
;-

;****************************************************
pro groom, energy, spectrum, nzero, line, back, integral
;
; ------------------ last changed: 07-apr-00 (aph)
;
; procedure by rk 4/99 to groom one spectrum.
; Outputs bkd-subtracted spectrum, background, and integral of line

;---------- fit first nzero+1 points are fit to a linear equation to get the background
line = POLY_FIT(energy(0:nzero), spectrum(0:nzero), 1 )
back = line(0) + energy*line(1)
spectrum = spectrum - back
; ----- report back info re relative amount of background
back = total(back)
integral=total(spectrum)
return
end
; ****************************************************

PRO stack_fit, axis=axis, verbose = verbose

@axis_com
@stack_process_com
COMMON volume_data, image_stack
@bsif_com

on_error,2


; ------------ read-in stack data -----------------------
stack = pickfile2(/READ, FILTER='*.ncb', /LPATH)
if strlen(stack) EQ 0 THEN RETURN  ; bail-out if no filename

stack_rb, stack		; read in stack in image_stack (Volume_data), energies in ev (stack_process_com)
 t=ax_name(stack)
 stackname_short = t(1)

; ------ read in reference spectra --------------------

;---- check if wish to use a fit parameter file -------------------
t = dialog_message('Read fit parameter file ?', /question)
if t(0) EQ 'Yes' then begin
    par_file = pickfile2(title = ' Select fit parameter file', $
           filter='*.par', /LPATH, DEFPATH=defpath )
	if strlen(par_file(0)) NE 0 then pars = ax_par_load(par_file) else goto, get_user_info
	if n_tags(pars) EQ 0 then begin
		goto, get_user_info
	endif else begin
		ncomp = pars.n
		comp_names = pars.names
; ---------- check if any have spaces - truncate at first space
		for i =  0, ncomp-1 do begin
			test= strpos(comp_names(i),' ')
			if test GT 0 then comp_names(i) = strmid(comp_names(i), 0, test)
			comp_names(i) = strtrim(comp_names(i),2)
			print, comp_names(i)
		endfor
		comp_files = pars.files
	endelse
endif else begin

; --- read in reference spectra files and component names one at a time
	get_user_info:
	if  keyword_set(AXIS) THEN BEGIN
		ncomp = get_num(Prompt= '# of components (1-8)', val=3, group = axis_id)
	endif else ncomp = get_num(Prompt= '# of components (1-8)', val=3)
	model = fltarr(ncomp, n_elements(ev))
	comp_names = strarr(ncomp)
	comp_files = strarr(ncomp)
	for i = 0, ncomp-1 do begin
		comp_files(i) = pickfile2(title = ' Spectrum of component '+ strtrim(string(i),2), $
		           filter='*.txt', /LPATH, DEFPATH=defpath )
		if comp_files(i) EQ '' then return  ; cancel at any time
		tmp = spc_load(file=comp_files(i))
	; --- establish name of component for file storage and header
		comp_names(i) = tmp.dl
		text = 'Name for component '+ strtrim(string(i),2)
		if keyword_set(axis) then begin
			comp_names(i) = get_text(prompt = text, val = comp_names(i),group = axis_id)
		endif else comp_names(i) = get_text(prompt = text, val = comp_names(i), group = axisID)
	 	test= strpos(comp_names(i),' ')
		if test GT 0 then comp_names(i) = strmid(comp_names(i), 0, test)
		comp_names(i) = strtrim(comp_names(i),2)
	endfor

; ------ save coefficient array for later use (optional)
	par_file = pickfile2(filter='*.par', /write, path = DefPath, $
	           title = 'Name of fit parameter file')
	if strlen(par_file(0)) NE 0 then ax_par_save,par_file,ncomp,comp_names,comp_files
ENDELSE

; check that range of reference data contains range of ev
; ----------- to write ---------------------


; check that the stack data is monotonic - if not, force monotonic
; ----- (aph; added 30-dec-03 - interpol failures may be cause of 'baulky' fits) -------------
d={x:ev,d:ev}
d = ax_sort_mono(d, /axis)
ev = d.x

; interpolate and store in AXIS data buffer
model = fltarr(ncomp, n_elements(ev))
for i = 0, ncomp-1 do begin
	tmp = spc_load(file=comp_files(i))
	tmp = ax_sort_mono(tmp,/axis)			; FORCE MONOTONIC !
	model(i,*)=interpol(tmp.d,tmp.x,ev)
	tmp = {t:'1d', d: model(i,*), x: ev, xl: 'X', dn:model(i,*), dl: 'interp '+ tmp.dl}
	CurBuf = i+1
	if  n_tags(tmp) NE 0 AND keyword_set(AXIS) THEN BEGIN
		HANDLE_VALUE, Data(Curbuf), tmp, /SET
		PlotBuf, Curbuf
	endif
endfor

; -------START OF KNEEDLER's FITSTACK procedure ---------

; determine number of components
print, ncomp, format="('Fitting spectrum at each pixel to ',i2,' components')"
; and # of energies (images)
n_E = n_elements(ev)

; groom reference spectra to have zero background (but they are
; assumed be normalized to each other in some way)
;if  keyword_set(AXIS) THEN BEGIN
;	nzero = get_num(Prompt='#bkground points for model spectra',Val = 0, group = axis_id)
;	nzero_img = get_num(Prompt='#bkground points for image spectra?', Val=0, group = axis_id)
;endif else begin
;	nzero = get_num(Prompt='#bkground points for model spectra',Val = 0)
;	nzero_img = get_num(Prompt='#bkground points for image spectra?', Val=0)
;endelse

nzero = 0
nzero_img = 0
Widget_control, /hourglass

; --------- dead code as no longer use groom (10-jun-04 aph) --------------
if nzero GT 0 then begin
	for i = 0, ncomp-1 do begin
		tmp = reform(model(i,*))
		groom, ev, tmp, nzero, line, bak, inte
		model(i,*) = tmp
		print,'reference spectrum ', strcompress(string(i+1)), $
		   '; integrated intensity = ',inte, ' : subtracted intensity =', bak
		print, '    Linear fit to background: slope = ', line(1), ' intercept = ', line(0)
	endfor
ENDIF

; switch plotting depending on keyword AXIS
if NOT keyword_set(axis) then begin
; plot reference signals
	window,0
	wset,0
	plot, ev,model(0,*)
	for i = 1, ncomp-1 do oplot,ev,model(i,*)
;		oplot, ev, 10*(line/max(line))
endif

; ------------ set up (x,y) scales from bsif_com information
x_step = (x_stop - x_start)/n_cols
y_step = (y_stop - y_start)/n_rows
img_x = x_start + findgen(n_cols)*x_step
img_y = y_start + findgen(n_rows)*x_step

; --------- set-up output arrays
const=fltarr(n_cols,n_rows)
comp = fltarr(ncomp, n_cols,n_rows)
chisq=const
bk=const
int=const
linear=const
if nzero_img GT 0 then begin  ; save the line fits to background
	slope = bk
	intercept = bk
endif
residual = float(image_stack)
if NOT keyword_set(axis) then begin
	window,1,xpos=0,ypos=200,xsize=n_cols,ysize=n_rows
	window,2,xpos=0,ypos=400,xsize=n_cols,ysize=n_rows
	if ncomp EQ 3 then window,3,xpos=0,ypos=600,xsize=n_cols,ysize=n_rows
	window,4,xpos=0,ypos=0,xsize=n_cols,ysize=n_rows
	loadct,0
endif

; set up regression fit params: x is array of model spectra ; -- removed(& line)
;  unweighted (weightsi=1, use /relative_weights keyword)
if ncomp GE 1 then tmp0 = transpose(reform(model(0,*)))
if ncomp GE 2 then tmp1 = transpose(reform(model(1,*)))
if ncomp GE 3 then tmp2 = transpose(reform(model(2,*)))
if ncomp GE 4 then tmp3 = transpose(reform(model(3,*)))
if ncomp GE 5 then tmp4 = transpose(reform(model(4,*)))
if ncomp GE 6 then tmp5 = transpose(reform(model(5,*)))
if ncomp GE 7 then tmp6 = transpose(reform(model(6,*)))
if ncomp GE 8 then tmp7 = transpose(reform(model(7,*)))
if ncomp EQ 1 then x=[tmp0]  ;,transpose(line)]
if ncomp EQ 2 then x=[tmp0,tmp1]   ; ,transpose(line)]
if ncomp EQ 3 then x=[tmp0,tmp1,tmp2] ;,transpose(line)]
if ncomp EQ 4 then x=[tmp0,tmp1,tmp2,tmp3]  ;,transpose(line)]
if ncomp EQ 5 then x=[tmp0,tmp1,tmp2,tmp3,tmp4]  ;,transpose(line)]
if ncomp EQ 6 then x=[tmp0,tmp1,tmp2,tmp3,tmp4,tmp5]  ;,transpose(line)]
if ncomp EQ 7 then x=[tmp0,tmp1,tmp2,tmp3,tmp4,tmp5,tmp6]  ;,transpose(line)]
if ncomp EQ 8 then x=[tmp0,tmp1,tmp2,tmp3,tmp4,tmp5,tmp6,tmp7]  ;,transpose(line)]
wts = make_array(n_E,value=1.0)
bak = 1.  & inte = 1.
for i=0,n_cols-1 do begin
  	for j=0,n_rows-1 do begin
		spectrum=reform(image_stack(i,j,*))
;		if nzero_img GT 0 then begin						; removed 10-jun-04 (aph)
;			groom, ev, spectrum, nzero_img, line, bak, inte
;			intercept(i,j) = line(0)
;			slope(i,j) = line(1)
;		endif
		bk(i,j)=bak
		int(i,j)=inte
		fit=regress(x, spectrum, wts, yfit, c, z, z, z, z, chi2, status, /relative_weight)
		for k = 0, ncomp -1 do begin
			comp(k,i,j)=fit(k)
		endfor
;			linear(i,j)=fit(ncomp)
		const(i,j)=c
		chisq(i,j)=chi2
		residual(i,j,*) = image_stack(i,j,*) - yfit
	;	wset,0
	;	plot,ev,spectrum
	;	oplot,ev,yfit
	;   print,fit,const(i,j),chisq(i,j),bk(i,j),int(i,j)
  	endfor

; ------- user feedback ... working ....
	if keyword_set(AXIS) THEN BEGIN
		WIDGET_CONTROL, Uprompt, SET_VALUE='STACK_FIT: line ' $
		    + strcompress(string(i)) +' of ' + strcompress(string(fix(n_cols)))
	endif


	if NOT keyword_set(axis) and keyword_set(verbose) then begin
;	NB the following ONLY works on graphics display devices with 3-channel (RBG) implemented
; plot results line-by-line   component 1 is red, 2 is green, 3 is blue
	  	top=max(comp)
		comp(*,0,0)=top

		if ncomp GE 1 then begin
			wset,1
			tvscl,reform(comp(0,*,*)),0,0,1
		endif
		if ncomp GE 2 then begin
			wset,2
			tvscl,reform(comp(1,*,*)),0,0,2
		endif
  		if ncomp GE 3 then begin
  			wset,3
  			tvscl,reform(comp(2,*,*)),0,0,3
  		endif
  		if ncomp GE 4 then begin
  			wset,4
  			tvscl,reform(comp(3,*,*)),0,0,3
  		endif
		if ncomp GE 5 then begin
  			wset,5
  			tvscl,reform(comp(4,*,*)),0,0,3
  		endif
  		if ncomp GE 6 then begin
  			wset,5
  			tvscl,reform(comp(5,*,*)),0,0,3
  		endif
  		if ncomp GE 7 then begin
  			wset,5
  			tvscl,reform(comp(6,*,*)),0,0,3
  		endif
  		if ncomp GE 8 then begin
  			wset,5
  			tvscl,reform(comp(7,*,*)),0,0,3
  		endif
; Attempt to generate 3-color composite
  		wset,ncomp+1
  		device, get_graphics_function= oldg,SET_GRAPHICS_FUNCTION = 7 ; set to "OR"
		if ncomp GE 1 then tvscl,reform(comp(0,*,*)),0,0,1
		if ncomp GE 2 then tvscl,reform(comp(1,*,*)),0,0,2
  		if ncomp ge 3 then tvscl,reform(comp(2,*,*)),0,0,3
  		device, set_graphics_function= oldg
	endif
; component 1 is green, component 2 is violet
;		comp1 = reform(comp(0,*,*))
;		comp2 = reform(comp(1,*,*))
;		c_max = max([comp1,comp2],min=c_min)
;	print, c_max, c_min
;		wset,1
;		tvscl,comp1        		; *256/c_max+c_min + 1  (failed attempt to give constant scale)
;		wset,2
;		tvscl,comp2,0,0,2        ; *256/c_max+c_min + 1

endfor

; --------- provide information about baclground fitting if done on data - removed 10-jun-04 (aph)
;if nzero_img GT 0 then begin
;	print, 'Background over image pixels'
;	print, 'Average slope = ', total(slope)/n_elements(slope), ' +/-', variance(slope)
;	print, 'Average intercept = ', total(intercept)/n_elements(intercept),' +/-', variance(intercept)
;endif

; finished the fit - store and display residuals
if keyword_set(axis) then begin

; ------------- define names for storing results
	t=ax_name(stack)
	storefile=t(1)
	fileroot = get_text(prompt = 'Root name for output files',val=storefile, group = axis_ID)
	fileroot = t(0) + fileroot + '_'

; -------- constant term in buffer 8----
; ----------- replace NaN values with 0.
	tmp = {t:'2d', d: const, x: img_x ,y: img_y, E: 0.0, xl:'X',yl: 'Y',  dl: 'constant '+ ' SF ' + stackname_short}
	bad_index = where(finite(tmp.d) EQ 0, count)
	if count GT 0 then begin
		tmp.d(bad_index) = 0.
		print, 'stack_fit: WARNING - NaN or Inf pixels set to 0.'
	endif
	HANDLE_VALUE, Data(8), tmp, /SET
	filename = fileroot+'constant.axb'
	overwrite_all = 0   ; reset overwrite flag in case was set by earlier uses
	test = axb_save(tmp,file=filename)
	PlotBuf,8

; -----------  store results in buffers and thumbnails
	for i = 0, ncomp-1 do begin
		tmp = {t:'2d', d: reform(comp(i,*,*)), x: img_x , E: 0.0, y: img_y, xl: 'X', $
		         yl: 'Y',  dl: comp_names(i)+ ' SF ' + stackname_short}
		bad_index = where(finite(tmp.d) EQ 0, count)
		if count GT 0 then begin
			tmp.d(bad_index) = 0.
			print, 'stack_fit: WARNING - NaN or Inf pixels set to 0.'
		endif
		tmp = ax_order(tmp)
		if ncomp EQ 4 then offset = 5 else offset = 4
		if ncomp GT 4 then offset = 1
		CurBuf = i + offset
		HANDLE_VALUE, Data(Curbuf), tmp, /SET
		Plotbuf, CurBuf
		filename = fileroot + comp_names(i) + '.axb'
		test = findfile(filename)
		if test(0) EQ filename AND overwrite_all NE 1 then begin
		    t = file_overwrite(group = axis_ID, val = 1)
		    if t EQ 0 then filename = pickfile2(/write, LFILE = filename)
		    if t EQ 2 then overwrite_all = 1
		endif
		test = axb_save(tmp,file=filename)
	endfor
	tmp = {t:'2d', d: chisq, x: img_x ,y: img_y, E: 0.0,  xl: 'X', yl: 'Y',  dl: 'residuals '+ ' SF ' + stackname_short}
	bad_index = where(finite(tmp.d) EQ 0, count)
	if count GT 0 then begin
		tmp.d(bad_index) = 0.
		print, 'stack_fit: WARNING - NaN or Inf pixels set to 0.'
	endif
	tmp = ax_order(tmp)
	HANDLE_VALUE, Data(9), tmp, /SET
	filename = fileroot + 'chi.axb'
	test = findfile(filename)
	if test(0) EQ filename AND overwrite_all NE 1 then begin
	    t = file_overwrite(group = axis_ID, val = 1)
	    if t EQ 0 then filename = pickfile2(/write, LFILE = filename)
	    if t EQ 2 then overwrite_all = 1
	endif
	test = axb_save(tmp,file=filename)
	CurBuf = 9
	PlotBuf,CurBuf
; ----------- store residual ?
   	if dialog_message('save residuals stack?',/question,/default_no, title = 'save residuals stack?') EQ 'Yes' then begin
		WIDGET_CONTROL, Uprompt, SET_VALUE='save residuals stack'
		filename = pickfile2(/WRITE, FILTER='*.ncb', /LPATH)
		if strlen(filename) GT 0 then begin
			image_stack = residual
			stack_wb, filename		; store for later viewing
		endif
	endif
; --------- store fitted background ?
	if nzero_img GT 0 then begin
		WIDGET_CONTROL, Uprompt, SET_VALUE='save background stack'
		filename = pickfile2(/WRITE, FILTER='*.ncb', /LPATH)
		if strlen(filename) GT 0 then begin
			for i = 0, n_E-1 do image_stack(*,*,i) = intercept + slope*eV(i)
			stack_wb, filename
		endif
	endif
endif
; ------------- display residuals
if NOT keyword_set(axis) then begin
	d_max=float(max(image_stack))
	window,5,xpos=200,ypos=200,xsize=n_cols,ysize=n_rows
	window,6,xpos=200,ypos=400,xsize=n_cols,ysize=n_rows
	for i=0,n_E-1 do begin
	;	print, ev(i)
		wset,ncomp + 2
	;	tv,500*residual,*,*,i)/image_stack(*,*,i)*256./d_max + 1
		tvscl,residual(*,*,i)     ; /image_stack(*,*,i)
		wset,ncomp + 3
	;	tv,image_stack(*,*,i)*256./d_max + 1
		tvscl,image_stack(*,*,i)
	;	wait,0.1
	endfor
endif
END


