;+
;
; Project   : s3drs reconstruction software
;                   
; Name      : (various test routines)
;               
; Purpose   : part of the s3drs test kit
;               
; Explanation: (archived here for completeness, not intended for users use)
;               
; Use       : (self-documented, see notes in code)
;    
; Written     : Sandy Antunes, NRC, 2005-2009
;               
;-            
; for this loop, 'scan' = just plots and shows the solution images

; quickcheck = 1 just plots log(renders) onscreen,
; quickcheck = 2 just plots log(img) onscreen,

; typical values for scaling:
;  model= any flux rope, scaling=0.01
;         any shell, scaling=5.0e-7
;         straight saito, scaling gets ignored
; raw minmaxes: saito=294,4.2e10; shellB=0.0,1.0; fluxy/fluxyB 0,5.6e5


; this loads a pixonmodel, adds a saito if asked, noises it, returns it

; choices are:
; 'saito','fluxrope','shell','fluxrope+saito','shell+saito','fluxy+saito'
;  and correspond to earlier 'dialme' choices of 'saito','fluxy','shellB',
;  and 'fluxyb'

FUNCTION howmodel,N2,modelstr,display=display,bright=bright,level=level
;print,'N2 is ',N2
  N=64; always make base model at N=64 then rebin up

  display=setyesno(display)
  bright=setyesno(bright)

  saito=0
  models=strsplit(modelstr,'\+',/extract)
  model=models[0]
  if (n_elements(models) gt 1) then begin
      if (models[1] eq 'saito') then saito=1
  end

  Nm=N/2
  stuff=[45.0,45.0]

  case model of
      'shell': begin
          offset=N/8
          shell=pixonmodels(Nm,'hshell2',ignore,stuff)
          timage=fltarr(N,N,N)
          timage[N/2-offset:(N/2)+Nm-1-offset,$
                 offset:offset+Nm-1,offset:offset+Nm-1]=shell
      end
      'farshell': begin
          offset=N/16
          shell=pixonmodels(Nm,'hshell2',ignore,stuff)
          timage=fltarr(N,N,N)
          timage[N/2-offset:(N/2)+Nm-1-offset,$
                 offset:offset+Nm-1,offset:offset+Nm-1]=shell
      end
      'fluxrope': begin
          timage=pixonmodels(N,'fluxrope',ignore,stuff)
          cuberot3d,timage,[0.0,20.0,-90.0],/reverse
          cuberot3d,timage,[90.0,0.0,0.0]
          cuberot3d,timage,[0.0,35.0,0.0]

      end
      'centerflux': begin
          timage=pixonmodels(N,'fluxrope',ignore,stuff)
      end
      'fluxy': begin
          timage=pixonmodels(N,'fluxrope',ignore,stuff)
          cuberot3d,timage,[30.0,0.0,0.0]
      end
      'saito': begin
          timage=fltarr(N,N,N); empty for now
          saito=1; add in 'saito' below
      end
      default: begin
      end
  endcase

;  help,timage
  if (bright) then timage=timage*100.0; for testing brighter events

  ; add to a saito background, if requested

  if (saito) then timage = timage+pixonmodels(N,'saito')
;  timage=timage + sqrt(timage)*randomn(seed,N,N,N)


  ; note since we are dealing with densities, not total electrons,
  ; we do not have to scale them by the voxel size.

  if (n_elements(level) ne 0) then timage=timage/level

  if (display) then threeview,fixed=196,/pixon,/log,timage

  timage2=rebin(timage,N2,N2,N2)

  return,timage2

END

; if /halo, rotates by an extra 90 degrees for fluxrope so it's a halo CME

; model choices are (as per 'dialme'):
;  fluxrope,fluxy,shell,fluxyB,shellB,shellB4,shellB8,shellB12,shellB16,saito
;  also need to give noiselevel (suggested: 1) and scaling (var.)
;
; adding /saitonoise adds saito to the noise level but _not_ the image.
;

; model choices are now: 'fluxrope','shell','farshell','saito'

; note we always use our latest 'local' Pixon environment here
;
; if you give the /compile option, it simply compiles this then
; exits, useful if you need the above model creation procedures.
;
; note tolerance sets opixon's reltol and abstol, default is 10e-3
;
; /spurious adds 2 bright single-pixel 'noise' events to the top of frame 1,
; just to see how Pixon reconciles irreconcilable input
;

PRO p_modeltests,model,N,gamma=gamma,theta=theta,halo=halo,show=show,$
                 noiselevel=noiselevel,scaling=scaling,noisy=noisy,$
                 nozero=nozero,sidemask=sidemask,nadjust=nadjust,$
                 xyzmask=xyzmask,zmask=zmask,compile=compile,mxit=mxit,$
                 bright=bright,nomasks=nomasks,nooccult=nooccult,$
                 saitonoise=saitonoise,LL=LL,megafile=megafile,level=level,$
                 polartrios=polartrios,polarpairs=polarpairs,nocut=nocut,$
                 seedshape=seedshape,timeless=timeless,tolerance=tolerance,$
                 loopfit=loopfit,cor1=cor1,spurious=spurious,$
                 testshift=testshift,recenter=recenter

compile=setyesno(compile)
if (compile) then return

on_error,2

pixon_env,'c',/local
print,'starting, ',systime()

halo=setyesno(halo)
cor1=setyesno(cor1)
cor2 = ~cor1; if not cor1, default is cor2
show=setyesno(show)
noisy=setyesno(noisy)
spurious=setyesno(spurious)
nozero=setyesno(nozero)
sidemask=setyesno(sidemask)
xyzmask=setyesno(xyzmask)
zmask=setyesno(zmask)
bright=setyesno(bright)
nomasks=setyesno(nomasks)
nocut=setyesno(nocut)
nooccult=setyesno(nooccult)
saitonoise=setyesno(saitonoise)
polartrios=setyesno(polartrios)
polarpairs=setyesno(polarpairs)
megafile=setyesno(megafile)
testshift=setyesno(testshift)
recenter=setyesno(recenter)
timeless=setyesno(timeless)
if (n_elements(loopfit) eq 0) then loopfit=0
if (n_elements(seedshape) eq 0) then seedshape='orthodox'
if (n_elements(noiselevel) eq 0) then noiselevel=1.0
if (n_elements(scaling) eq 0) then scaling=1.0

if (n_elements(N) eq 0) then N=128

if (n_elements(model) eq 0) then model='saito'; or 'fluxy' or 'shell' or 'farshell'

; default is x,y,z
if (n_elements(gamma) eq 0) then theta=[0.0,0.0,90.0]
if (n_elements(gamma) eq 0) then gamma=[0.0,90.0,0.0]
; default is planar if theta not given
if (n_elements(theta) eq 0) then theta=gamma*0.0

; create optional polarity details
if (polartrios or polarpairs) then begin
    iorig=n_elements(gamma)
    if (polartrios) then gamma=[gamma,gamma,gamma] else gamma=[gamma,gamma]
    if (polartrios) then theta=[theta,theta,theta] else theta=[theta,theta]
    polar=make_array(n_elements(gamma),/string,value='p')
    for i=0,iorig-1 do begin
        if (polartrios) then polar[i]='0' else polar[i]='t'
        if (polartrios) then polar[i+iorig]='120' else polar[i+iorig]='r'
        if (polartrios) then polar[i+2*iorig]='240'
    end
end

ict=n_elements(gamma); note we tally this before the additional masks

if (zmask) then gamma=[gamma,0.0]
if (xyzmask) then gamma=[gamma,0.0,90.0,0.0]
if (zmask) then theta=[theta,90.0]
if (xyzmask) then theta=[theta,90.0,0.0,0.0]

if (halo) then gamma=gamma+90.0

; input 'model' can be a string or an actual model
if (n_elements(model) gt 1) then timage=model else $
  timage=howmodel(N,model,bright=bright,level=level)      ;,/display)
if (n_elements(model) gt 1) then modelname='custom' else modelname=model

;dialme,model,scaling,timage,noiselevel=1.0,/nosave
;timage=rebin(timage,N,N,N)

; hack needed due to some weird input problem

ie=where (timage le 0)
if (ie[0] ne -1) then timage[ie]=0.0

; then make data:
prep_3drv,pxwrapper,N=N,gamma=gamma,theta=theta,cor2=cor2,cor1=cor1,$
  polar=polar,LL=LL,ftol=ftol

  ; store the specifics of this setup, for record-keeping
  if (xyzmask) then pxwrapper.sc.xyzmask=1
  if (zmask) then pxwrapper.sc.zmask=1
  if (sidemask) then pxwrapper.sc.sidemask=1
  if (nozero) then pxwrapper.sc.nozero=1
;  if (excerpt) then pxwrapper.sc.excerpt=1
;  if (occultleft or occultright or occulttop or occultbottom) then $
;    pxwrapper.sc.occult=1

mmon=  int2str(get_month(strmid(systime(),4,3)),2)
dday=  int2str(long(strmid(systime(),8,2)),2)
hhr=   strmid(systime(),22)
datestamp=mmon+dday+hhr
if (timeless) then datestamp='auto'
pxwrapper.sc.gname=modelname+strtrim(string(N),2)+'_'+$
  strjoin(strtrim(string(floor(gamma)),2),',')+'_'+datestamp
;+'_noise'+$
;  strtrim(string(noiselevel),2)+'_mscale'+$
;  strtrim(string(scaling),2)+'_'

if (bright) then pxwrapper.sc.gname = 'bright'+pxwrapper.sc.gname

if (zmask) then pxwrapper.sc.gname += '_zmask'
if (xyzmask) then pxwrapper.sc.gname += '_xyzmask'
if (sidemask) then pxwrapper.sc.gname += '_sidemask'
if (nozero) then pxwrapper.sc.gname += '_nozero'
if (nooccult) then pxwrapper.sc.gname += '_nooccult'
if (spurious) then pxwrapper.sc.gname += '_spurious'
if (testshift) then pxwrapper.sc.gname += '_testshift'
if (recenter) then pxwrapper.sc.gname += '_recenter'
pxwrapper.sc.gname += '_starter='+seedshape

print,'Setting up and solving: ',pxwrapper.sc.gname

if (N gt 128) then pxwrapper.sp=p_tweak_kernel(pxwrapper.sp,2,1)

datum=fltarr(N,N,ict)
for i=0,ict-1 do datum[*,*,i]=pr_render(pxwrapper,timage,i)


if (testshift) then begin
; Test here of image shifting.  This shows that the images do not
; have to be co-aligned, rather you just have to always specify the
; sun center accurately within each 2D data frame.
    pxwrapper.sr.sunk0[0]=ptr_new([N/4,31.5])
    pxwrapper.sr.k0[0]=ptr_new([N/4,31.5])
    datum[31.5,31.5,0]=max(datum)
    tv_multi,fixed=256,/log,datum,header='original input data plus bright dot'
    dtemp=datum[*,*,0]*0.0
    dtemp[0:N/2,*]=datum[(N/4)-1:(3*N/4)-1,*,0]
;    tv_multi,fixed=256,/log,dtemp
    dtemp[N/4,31.5]=0
;    tv_multi,fixed=256,/log,dtemp
    datum[*,*,0]=dtemp
    tv_multi,fixed=256,/log,datum,header='shifted input data, if bright single pixel dot then debug in p_modeltests.pro'
end

; noise is done by converting MSB to DN, then using our
; Cor2A DN->photons+exptime factor, then converting back,
; all assuming a usual 6-second frame.


;;;if (n_elements(nadjust) eq 0) then nfactor=0.01 else nfactor=nadjust; 6.0
;if (n_elements(nadjust) eq 0) then nfactor=1.0 else nfactor=nadjust;
if (n_elements(nadjust) eq 0) then begin
    if (N eq 64) then nfactor=10.0 else if (N eq 128) then nfactor=100.0 else nfactor=1.0
end else begin
    nfactor=nadjust             ; 6.0
end
pxwrapper.sc.nadjust=nfactor

if (spurious) then begin
    ; adds 2 bright single pixels to image
    datum[N/4,(3*N/4)+2,0]=max(datum)
    datum[3*N/4,(3*N/4)-2,0]=max(datum)
;    tv_multi,/log,datum
end

if (saitonoise) then addsaito=howmodel(N,'saito',level=level) else addsaito=0.0
;noise=p_quicknoise(datum+addsaito,pxwrapper=pxwrapper)/nfactor
noise=p_simplenoise(datum+addsaito,pxwrapper)/nfactor

; add noise to our perfect data to 'mix it up a bit'
if (noisy) then begin
    flatfield = max(datum)*1e-2
    datum=abs(datum +noise*randomn(seed,N,N,N) +flatfield*randomn(seed,N,N,N))
  ; and make sure we now calculate the real noise for our now-noisy data
;  noise=p_quicknoise(datum,pxwrapper=pxwrapper)/nfactor
  noise=p_simplenoise(datum+addsaito,pxwrapper)/nfactor
end

;dn2msb=2.7E-12
;datum2= datum / dn2msb
; now we can use our 'quickcme' noise formulation
;
;Nfactor=2048/float(N)  ; e.g. 1/16 for N=128 so Nfactor=16
;const=1.0 / (sqrt(90.0) * Nfactor)
; 100000 factor is a HACK!!!!
;noise2=const * sqrt(datum2) * 100000.0
;noise = noise2 * dn2msb


ie = where (datum gt 0)
if (ie[0] eq -1) then dmin=0.0 else dmin=min(datum[ie])

;dmin=max(datum)*1e-5
;print,max(datum),dmin

if (recenter) then p_recenter,pxwrapper,datum,noise,/show


masks=datum*0.0
if (nooccult) then begin
    masks=masks + 1.0; flat everywhere
end else begin
    masks=p_getmask(sr=pxwrapper.sr)
end
if (nomasks) then masks=masks*0.0 + 1.0

;tv_multi,/log,fixed=128,datum*masks

if (sidemask) then begin
    ie=where(masks eq 0)
    if (ie[0] ne -1) then datum[ie]=dmin*0.1
    if (ie[0] ne -1) then noise[ie]=dmin*10.0
    masks=masks*0.0+1.0; masks are no longer relevant as data was altered
end

if (zmask or xyzmask) then begin
;    blankZ=p_getmask('COR2',N); includes center occulted
    if (cor1) then det='COR1' else det='COR2'
    blank=p_getmask(det,N,rocc=-1); just edge occulted
;    datumblank=blank*0.0 + dmin*0.1 ; low everywhere
    datumblank= blank*0.0;   zero everywhere we care
    maskZ= ~blank   ;Z
    maskXY= ~blank
;    sigmablank= datumblank * 10         ; wildly noisy
    sigmablank=datumblank + max(datum)*1E-7; make noise 7 oom below peak signal
tv_multi,maskZ,header='mask'
tv_multi,/log,/show,datumblank,header='data'
tv_multi,/log,/show,sigmablank,header='sigma'

    if (zmask) then datum=[[[datum]],[[datumblank]]]
    if (xyzmask) then $
      datum=[[[datum]],[[datumblank]],[[datumblank]],[[datumblank]]]
    if (zmask) then noise=[[[noise]],[[sigmablank]]]
    if (xyzmask) then $
      noise=[[[noise]],[[sigmablank]],[[sigmablank]],[[sigmablank]]]
    masks[*,*,ict]=maskZ
    if (zmask) then pxwrapper.sr.detector[ict]='z mask'
    if (xyzmask) then begin
        masks[*,*,ict]=maskXY
        pxwrapper.sr.detector[ict]='z mask'
        masks[*,*,ict+1]=maskXY
        pxwrapper.sr.detector[ict+1]='y mask'
        masks[*,*,ict+2]=maskXY
        pxwrapper.sr.detector[ict+1]='x mask'
    end
end

if (nozero) then p_nozero,datum,noise

; *** false zero mod, to remove actual data zeroes
;if (nozero) then begin
;    dmin=min(datum)
;    ie=where(datum le 0)
;    if (ie[0] ne -1) then datum[ie]=dmin*0.1
;end

; check S/N

sn=ratio(datum*masks,noise)

save,datum,noise,pxwrapper,timage,file=pxwrapper.sc.gname+'_datum.sav'
;pxwrapper.sr.normd=1e-6

print,'Model is ',pxwrapper.sr.imodel

if (n_elements(tolerance) ne 0) then begin
    pxwrapper.sp.reltol=tolerance
    pxwrapper.sp.abstol=tolerance
    ; both get applied in p_osetup, or later you can set it manually with
    ; opixon->set,abstol=pxwrapper.sp.abstol,reltol=pxwrapper.sp.reltol
end


p_generate_pixon,'classic',pxwrapper,oPixon
;help,pxwrapper.sr.nd,datum,noise,masks
p_insert_data,pxwrapper,datum,oPixon,sigma=noise,masks=masks

orthodox=0
case seedshape of
    'jumpstart': begin
        startimage=simple3drv(datum,pxwrapper=pxwrapper)
        startimage = startimage*max(datum)/max(pr_render(pxwrapper,startimage))
    end
    'biginitfov': begin
        startimage=pr_fov(pxwrapper)
    end
    'medinitfov': begin
        startimage=pr_fov(pxwrapper)/(float(N)^3)
    end
    'initfov': begin
        startimage=pr_fov(pxwrapper)
        startimage=startimage / (max(startimage))
    end
    'oddinitfov': begin
        startimage=pr_fov(pxwrapper)
        startimage=startimage / (max(startimage))
        startimage = startimage * float(N)^3
    end
    'perfect': begin
        startimage=timage
    end
    'flat': begin
        startimage=fltarr(N,N,N) + 1.0/float(N)
    end
    'bigflat': begin
        startimage=fltarr(N,N,N) + 1.0
    end
    'biggerflat': begin
        startimage=fltarr(N,N,N) + 1E5
    end
    'zero': begin
        startimage=fltarr(N,N,N)*0.0
    end
    'tinycube': begin
        startimage=fltarr(N,N,N)+1.0
        startimage[0:N/4,*,*]=0
        startimage[3*N/4:N-1,*,*]=0
        startimage[*,0:N/4,*]=0
        startimage[*,3*N/4:N-1,*]=0
        startimage[*,*,0:N/4]=0
        startimage[*,*,3*N/4:N-1]=0
    end
    'bigtinycube': begin
        startimage=fltarr(N,N,N)+1E5
        startimage[0:N/4,*,*]=0
        startimage[3*N/4:N-1,*,*]=0
        startimage[*,0:N/4,*]=0
        startimage[*,3*N/4:N-1,*]=0
        startimage[*,*,0:N/4]=0
        startimage[*,*,3*N/4:N-1]=0
    end
    else: begin
;    oPixon->replace,image=fltarr(N,N,N)+1.0/float(N^3)
        pxwrapper.sc.gname=pxwrapper.sc.gname+'_orthodox'
        orthodox=1
    end
endcase

if (orthodox eq 0) then begin
    print,total(startimage),mean(startimage),pxwrapper.sc.gname
    oPixon->replace,image=startimage
end

; check it
if (show) then begin
    masks2=wt2mask(oPixon)
    datum2=stereo_fetch(pxwrapper,oPixon,/dat)
    sigma2=stereo_fetch(pxwrapper,oPixon,/sigma)
    nframes=n_elements(masks2[0,0,*])
    lbase=['masks out','data out','sigma out','sigma in','sn > 1']
    lset=['']
    for i=0,n_elements(lbase)-1 do $
      lset=[lset,make_array(/string,nframes,value=lbase[i])]
    lset=lset[1:*]
    dset=[[[masks2]],[[datum2]],[[sigma2]],[[noise]],[[sn > 1]]]
    tv_multi,dset,line=nframes,/local,fixed=96,/log,/show,title='in&out',$
      labels=lset,oom=2
end

odata=opixon->get(/data,/nop)
osigma=opixon->get(/sigma,/nop)
;print,'limits:'
;print,minmax(datum),minmax(noise),minmax(odata),minmax(osigma)


;save,oPixon,datum,noise,pxwrapper,timage,$
;  file=modelname+strtrim(string(N),2)+'_opixon.sav'

print,'fit starting, ',systime()
if (n_elements(mxit) ne 0) then opixon->set,mxit=mxit

if (orthodox) then begin
    oPixon->full,/restart
end else if (loopfit eq 1) then begin
    pxwrapper.sc.gname += '_loopfit1'
    for i=0,2 do begin
        oPixon->Pixon::NNLS
        oPixon->Map
        oPixon->Pseudo
    end
    ; no improvement over existing pixon full, just takes longer/retries more
end else if (loopfit eq 2) then begin
    pxwrapper.sc.gname += '_loopfit2'
    oPixon->Pseudo
    for i=0,2 do begin
        oPixon->Map
        oPixon->PixonSmooth,/image
        oPixon->Pseudo
    end
    ; no improvement over existing pixon full, just takes longer/retries more
end else if (loopfit eq 3) then begin
    pxwrapper.sc.gname += '_loopfit3'
    for i=0,2 do begin
        oPixon->Pseudo
        oPixon->Map
        oPixon->PixonSmooth,/pxnidx; ends up resetting Q at start of each loop
    end
    oPixon->Pseudo ; in the end, fit was worse
end else begin
    oPixon->full
end
;opixon->free,/phi,/pxnidx;  this really screws things up, no first loop!?!?!
;oPixon->full

print,'fit ended, ',systime()

fname=pxwrapper.sc.gname+'_solved.sav'

p_update_runstats,pxwrapper,oPixon

save,pxwrapper,oPixon,timage,file=fname
print,'Saved ',fname

if (megafile) then begin
  sh=0
  fl=0
  if (model eq 'shell' or model eq 'farshell') then sh=1 else $
    if (model eq 'fluxrope') then fl=1
  best=~(sh+fl)

  p_standardplots,pxwrapper,oPixon,synthetic=timage,/mega,shellbest=sh,$
    fluxbest=fl,best=best,mxit=150,nocut=nocut

end

;p_standardplots,pxwrapper,oPixon,/cap,oom=2,/noscreen,synthetic=timage

opixon->free

if (show) then begin
  img=stereo_fetch(pxwrapper,opixon,/im)
  threeview,/pixon,fixed=128,img,/log,oom=2,title='Image out'
  qmap=stereo_fetch(pxwrapper,oPixon,/qmap)
  tv_multi,fixed=128,qmap,title='Q Map',/log,oom=2
end

print,'done, ',systime()

END

pro howdiag,pxwrapper,oPixon

  tv_multi,fixed=128,/log,stereo_fetch(pxwrapper,oPixon,/data),title='data'
  tv_multi,fixed=128,/log,stereo_fetch(pxwrapper,oPixon,/sigma),title='sigma'
  tv_multi,fixed=128,/log,stereo_fetch(pxwrapper,oPixon,/render),title='render'
  threeview,/pixon,fixed=128,/log,$
    stereo_fetch(pxwrapper,oPixon,/image),title='sol'


end