#!/usr/bin/env python

# coding: utf-8

# # Test the one-line CLEDB inversion on data.

# The script version "test_1line.py" is generated by exporting this notebook via file -> export as -> executable script. Do not directly edit that version!

# In[ ]:

## Inversion Constants and control parameter imports.

import constants as consts

import ctrlparams

params=ctrlparams.ctrlparams() ##Initialize and use a shorter label

## Needed modules

import importlib ## Reloads imports (python default settings will not reload small changes).

import numpy as np ## Defines variables to store/load headers

from astropy.io import fits ## To properly load DKISt Cryo-NIRSP data

# ### 1. Import the Fe XIII 1074 nm and 1079 nm synthetic CLE/MURAM or observed MLSO CoMP or DKIST Cryo-NIRSP data

# #### 1.a CLE 3dipole along LOS simulation example

# In[ ]:

# ## observations of a 3 dipole coronal structure of a Fe XIII combined observation

# ## sobs1-3 are the independent dipoles

# ## sobsa is the combined 3 dipole output

# ## waveA and waveB are the wavelength arrays for the two Fe XIII lines

# sobs1,sobs2,sobs3,sobsa,waveA,waveB = np.load('obsstokes_3dipole_hires_fullspectra_1kg.pkl', allow_pickle=True) ## Generally returns a 1 element list, unpack via [0]

# ### reversing of the wavelength range. THIS IS NEEDED! CLE writes frequency-wise, so wavelengths are reversed in the original datacubes!!!!!!

# sobsa=sobsa[:,:,::-1,:]

# #waveA=waveA[::-1] ##the wave arrays are not needed by the inversion. the information is reconstructed from keywords

# #waveB=waveB[::-1]

# if params.integrated == False:

# sobs_in=[sobsa[:,:,:,0:4]] ## The simulated data is a x*y*wl*8 dimensional array with both lines included. ## create a 2line standard input for an actual observation.

# ## A fake minimal header for CLE

# ## This assumes the reference pixel is in the left bottom of the array; the values are in solar coordinates at crpixn in r_sun. The wavelength raferences (vacuum), ranges and spectral resolutions are unique. See CLE outfiles and grid.dat.

# ## sobsa.shape == (600, 280, 30, 8)

# head_in =[{'CRPIX1':0, 'CRPIX2':0, 'CRPIX3': 14, 'CRVAL1': -0.75, 'CRVAL2': 0.8, 'CRVAL3': 1074.61446, 'CDELT1': (0.75-(-0.75))/600, 'CDELT2': (1.5-0.8)/280, 'CDELT3': 0.0247, 'LINEWAV': "1074", 'INSTRUME': "CLE-SIM"}]

# ## Delete the large arrays from memory

# sobsa = 0

# #### 2.a MURAM data example loading

# In[ ]:

# ## load the fake observation muram data

# ## FE XIII 1074+1079

# # with open('obsstokes_avg_muram3.pkl','rb') as f:

# # sobsa = pickle.load(f)

# sobsa=np.load('obsstokes_avg_muram3.pkl', allow_pickle=True)[0] ## Generally returns a 1 element list, unpack via [0]

# if params.integrated == False:

# sobs_in=[sobsa[:,:,:,0:4]] ## The simulated data is a x*y*wl*8 dimensional array with both lines included. ## create a 2line standard input for an actual observation.

# # ## A fake minimal header for MURAM

# # ## This assumes the reference pixel is in the left bottom of the array; the values are in solar coordinates at crpixn in r_sun (from muram xvec and yvec arrays). The wavelength raferences (vacuum), ranges and spectral resolutions are unique (muram wvvec1 and wvvce2 arrays).

# head_in =[{'CRPIX1':0, 'CRPIX2':0, 'CRPIX3':0, 'CRVAL1': -0.071, 'CRVAL2': 0.989, 'CRVAL3': 1074.257137, 'CDELT1': 0.0001379, 'CDELT2': 0.0000689, 'CDELT3': 0.0071641, 'LINEWAV': "1074", 'INSTRUME': "MUR-SIM"}]

# ## Delete the large arrays from memory

# del sobsa

# else:

# print("Control parameter set to integrated. Data will not be read properly downstream!")

# #### 3.a DKIST Cryo-NIRSP Provisional example data (Integrated required).

# In[ ]:

# import dkist ## Optionally needed only if loading DKIST data

# if params.integrated == True:

# ##Loads the observation metadata into memory. Unique datasets are preserved here.

# asdf_file_1 = './obsstokes_cryonirsp_L1_BZJOM_metadata.asdf' ## raster dataset - 1074

# head_in = [dkist.load_dataset(asdf_file_1).headers] ### list of two headers

# sobsa = np.zeros((head_in[0][0]['CNNUMSCN'],head_in[0][0]['DNAXIS2'],4),dtype=np.float32) ## sobsa will not have a wavelength point dimension in this case

# sobsa = np.load('./obsstokes_cryonirsp_BZJOM_int1074') ## 1074 fit data

# sobs_in=[sobsa] ## The Xryo-NIRSP data is a x*y*wl*4 dimensional array with one lines included. ## create a 2line standard input for an actual observation.

# ## Delete the large arrays from memory

# sobsa = 0

# MLSO CoMP/uCoMP data can not be used here due to the absence of Stokes V.

# ### 2. Test the CLEDB_PREPINV module with synthetic data.

# ##### Remember to set your personal options and database paths in the ctrlparams class (in the parent directory) before continuing.

# In[ ]:

import CLEDB_PREPINV.CLEDB_PREPINV as prepinv ##imports from the CLEDB_PREPINV subdirectory

# ##### Preprocess the observation "files".

# In[ ]:

importlib.reload(prepinv) ## If module is modified, reload the contents

sobs_tot,yobs,snr,background,issuemask,wlarr,keyvals,sobs_totrot,aobs,dobs=prepinv.sobs_preprocess(sobs_in,head_in,params)

# ##### Databases are not required in the one line case.

# #### At this point, all necessary data are loaded into memory for fast processing.

# ### 3. Test the CLEDB_PROC module with the same synthetic data.

# In[ ]:

import CLEDB_PROC.CLEDB_PROC as procinv

# ##### Process the spectroscopy outputs

# In[ ]:

importlib.reload(procinv) ## If module is modified, reload the contents

specout = procinv.spectro_proc(sobs_in,sobs_tot,snr,issuemask,background,wlarr,keyvals,consts,params)

# ##### Process the LOS magnetic fields from the requested line

# In[ ]:

importlib.reload(procinv) ## If module is modified, reload the contents

blosout=procinv.blos_proc(sobs_tot,snr,issuemask,keyvals,consts,params)

# ### All should be good if we reached this point; all the outputs should be computed and saved in memory.

# ## 4. OPTIONAL tidbits

# In[ ]:

##optionally needed libraries and functions

from datetime import datetime

import os

import glob

import pickle

import numpy as np

from matplotlib import pyplot as plt

## interactive plotting; use only on local machines if widget is installed

get_ipython().run_line_magic('matplotlib', 'widget')

# colorbar function to have nice colorbars in figures with title

def colorbar(mappable,*args,**kwargs):

from mpl_toolkits.axes_grid1 import make_axes_locatable

import matplotlib.pyplot as plt

last_axes = plt.gca()

ax = mappable.axes

fig = ax.figure

divider = make_axes_locatable(ax)

cax = divider.append_axes("right", size="5%", pad=0.05)

cbar = fig.colorbar(mappable, cax=cax,format='%.3f')

#cbar.formatter.set_powerlimits((0, 0))

title= kwargs.get('title', None)

cbar.set_label(title)

plt.sca(last_axes)

return cbar

# ### 4.a DUMP results (optional)

# In[ ]:

## Remove old file saves and keep just the last run

lst=glob.glob('./outparams_2line*.pkl')

if len(lst) >0:

for i in range(len(lst)):

os.remove(lst[i])

## save the last run

datestamp = datetime.now().strftime("%Y%m%d-%H:%M:%S")

# if not os.path.exists('./testrun_outputs'): ## make an output directory to keep things clean

# os.makedirs('./testrun_outputs')

# if os.path.isfile(f"testrun_outputs/outparams_2line_{datestamp}.pkl"):

# print("Save data exists; Are you sure you want to overwrite?")

# else:

# if params.iqud == False: ## all outputs saved to disk

# with open(f'./testrun_outputs/outparams_2line_{datestamp}.pkl', 'wb') as f: # Python 3: open(..., 'wb')

# pickle.dump([specout,blosout,invout,sfound], f)

# else: ## no specout and blosout for iqud

# with open(f'./testrun_outputs/outparams_2line_{datestamp}.pkl', 'wb') as f: # Python 3: open(..., 'wb')

# pickle.dump([invout,sfound], f)

# ### 4.b PLOT the outputs (optional)

# In[ ]:

## Plot bounds and ranges utils

linen=0 ## choose which line to plot; range is [0:1] for 2 line input

## Plot subrange definition for some example input observations:

## !!!These bounding ranges will have to be calculated for particulare observations fed!!!

instrument = head_in[0]['INSTRUME'].strip()

if instrument == "CLE-SIM": ## 3dipole simulation example (1.a)

srx1=230

srx2=400

sry1=65

sry2=195

rnge=[0.8,1.5,-1.1,1.1]

elif instrument == "MUR-SIM": ## muram simulation example (2.a)

srx1=0

srx2=1024

sry1=0

sry2=1024

rnge=[0.989,1.060,-0.071,0.071]

elif instrument == "COMP": ## CoMP observation example (4.a)

srx1=0

srx2=620

sry1=0

sry2=620

rnge=[-1.384,1.384,-1.384,1.384]

elif instrument == "UCoMP": ## UCoMP observation example (4.b)

srx1=0

srx2=1024

sry1=0

sry2=1280

rnge=[-1.942,1.942,-1.553,1.553]

elif instrument == "CRYO-NIRSP": ## Cryo-NIRSP observation example (3.a)

srx1=0

srx2=51

sry1=0

sry2=1917

rnge=[-500,250,870,1200]

else:

raise ValueError(f"Unknown example: {instrument}")

print(f"Loaded bounds for observation example: {instrument}")

# ### Plot spectroscopy outputs

# In[ ]:

if params.integrated != True:

fig, plots = plt.subplots(nrows=3, ncols=4, figsize=(12,10))

## remove the 0 values and unreasonable/outlier values outside of the range of the four lines.

mx = np.ma.masked_array(specout[srx1:srx2,sry1:sry2,linen,0], mask=specout[srx1:srx2,sry1:sry2,linen,0]==0)

mx = np.ma.masked_array(mx, mask=specout[srx1:srx2,sry1:sry2,linen,0]>= 3950)

mx = np.ma.masked_array(mx, mask=specout[srx1:srx2,sry1:sry2,linen,0]<= 1070)

vvmin=np.min(mx)

vvmax=np.max(mx)

ab=plots[0,0].imshow(specout[srx1:srx2,sry1:sry2,linen,0],extent=rnge,origin='lower',vmin=vvmin,vmax=vvmax,cmap='PuOr',interpolation='none')

plots[0,0].set_title('Wavelength')

colorbar(ab,title="[nm]")

plots[0,0].set_ylabel('Z [R$_\odot$]')

#plots[0,0].set_xlabel('Y [R$_\odot$]')

############################################################

## for correctly scaling in doppler scales # remove unreasonable/outlier values of > 2 nm.

mx = np.ma.masked_array(specout[srx1:srx2,sry1:sry2,linen,1], mask=specout[srx1:srx2,sry1:sry2,linen,1]>= 2 )

mx = np.ma.masked_array(mx, mask=specout[srx1:srx2,sry1:sry2,linen,1]<= -2 )

vvmax=np.max(mx)

vvmin=np.min(mx)

if np.abs(vvmin) > np.abs(vvmax):

vr=np.abs(vvmin)

else:

vr=np.abs(vvmax)

ab=plots[0,1].imshow(specout[srx1:srx2,sry1:sry2,linen,1],extent=rnge,origin='lower',vmin=-vr,vmax=vr,cmap='seismic',interpolation='none')

plots[0,1].set_title('Doppler shift')

colorbar(ab,title="[nm]")

#plots[0,1].set_ylabel('Z [R$_\odot$]')

#plots[0,1].set_xlabel('Y [R$_\odot$]')

############################################################

## for correctly scaling in speed scales # remove unreasonable/outlier values of > 1000km/s.

mx = np.ma.masked_array(specout[srx1:srx2,sry1:sry2,linen,2], mask=specout[srx1:srx2,sry1:sry2,linen,2]>= 1000 )

mx = np.ma.masked_array(mx, mask=specout[srx1:srx2,sry1:sry2,linen,2]<= -1000 )

vvmax=np.max(mx)

vvmin=np.min(mx)

if np.abs(vvmin) > np.abs(vvmax):

vr=np.abs(vvmin)

else:

vr=np.abs(vvmax)

## back to plotting

ab=plots[0,2].imshow(specout[srx1:srx2,sry1:sry2,linen,2],extent=rnge,origin='lower',vmin=-vr,vmax=vr,cmap='seismic',interpolation='none')

plots[0,2].set_title('Doppler shift')

colorbar(ab,title="[km s$^{-1}$]")

#plots[0,2].set_ylabel('Z [R$_\odot$]')

#plots[0,2].set_xlabel('Y [R$_\odot$]')

############################################################

ab=plots[0,3].imshow(specout[srx1:srx2,sry1:sry2,linen,7],extent=rnge,origin='lower',vmin=0,vmax=np.max(specout[srx1:srx2,sry1:sry2,linen,7]),cmap='Reds',interpolation='none')

plots[0,3].set_title('Stokes I bkg.')

colorbar(ab,title="Signal [erg cm$^{-2}$ s$^{-1}$]")

#plots[0,3].set_ylabel('Z [R$_\odot$]')

#plots[0,3].set_xlabel('Y [R$_\odot$]')

############################################################

ab=plots[1,0].imshow(specout[srx1:srx2,sry1:sry2,linen,3],extent=rnge,origin='lower',vmin=0,vmax=np.max(specout[srx1:srx2,sry1:sry2,linen,3]),cmap='Reds',interpolation='none')

plots[1,0].set_title('Stokes I int. (linecore-bkg)')

colorbar(ab,title="Signal [erg cm${-2}$ s$^{-1}$]")

plots[1,0].set_ylabel('Z [R$_\odot$]')

#plots[1,0].set_xlabel('Y [R$_\odot$]')

############################################################

vvmin=np.min(specout[srx1:srx2,sry1:sry2,linen,4])

vvmax=np.max(specout[srx1:srx2,sry1:sry2,linen,4])

if np.abs(vvmin) > np.abs(vvmax):

vr=np.abs(vvmin)

else:

vr=np.abs(vvmax)

ab=plots[1,1].imshow(specout[srx1:srx2,sry1:sry2,linen,4],extent=rnge,origin='lower',vmin=-vr,vmax=vr,cmap='seismic',interpolation='none')

plots[1,1].set_title('Stokes Q int. (linecore-bkg)')

colorbar(ab,title="Signal [erg cm$^{-2}$ s$^{-1}$]")

#plots[1,1].set_ylabel('Z [R$_\odot$]')

#plots[1,1].set_xlabel('Y [R$_\odot$]')

############################################################

vvmin=np.min(specout[srx1:srx2,sry1:sry2,linen,5])

vvmax=np.max(specout[srx1:srx2,sry1:sry2,linen,5])

if np.abs(vvmin) > np.abs(vvmax):

vr=np.abs(vvmin)

else:

vr=np.abs(vvmax)

ab=plots[1,2].imshow(specout[srx1:srx2,sry1:sry2,linen,5],extent=rnge,origin='lower',vmin=-vr,vmax=vr,cmap='seismic',interpolation='none')

plots[1,2].set_title('Stokes U int. (linecore-bkg)')

colorbar(ab,title="Signal [erg cm$^{-2}$ s$^{-1}$]")

#plots[1,2].set_ylabel('Z [R$_\odot$]')

#plots[1,2].set_xlabel('Y [R$_\odot$]')

############################################################

vvmin=np.min(specout[srx1:srx2,sry1:sry2,linen,6])

vvmax=np.max(specout[srx1:srx2,sry1:sry2,linen,6])

if np.abs(vvmin) > np.abs(vvmax):

vr=np.abs(vvmin)

else:

vr=np.abs(vvmax)

ab=plots[1,3].imshow(specout[srx1:srx2,sry1:sry2,linen,6],extent=rnge,origin='lower',vmin=-vr,vmax=vr,cmap='seismic',interpolation='none')

plots[1,3].set_title('Stokes V int. (linecore-bkg)')

colorbar(ab,title="Signal [erg cm$^{-2}$ s$^{-1}$ nm$^{-1}$]")

#plots[1,3].set_ylabel('Z [R$_\odot$]')

#plots[1,3].set_xlabel('Y [R$_\odot$]')

############################################################

## remove unreasonable/outlier values from plotting range. Typical Fe XIII widths are 0.15. Outliers <0.5 are excluded from scaling.

mx = np.ma.masked_array(specout[srx1:srx2,sry1:sry2,linen,8], mask=specout[srx1:srx2,sry1:sry2,linen,8] >= 0.50 )

vvmax=np.max(mx)

ab=plots[2,0].imshow(specout[srx1:srx2,sry1:sry2,linen,8],extent=rnge,origin='lower',vmin=0.0,vmax=vvmax,cmap='YlGnBu',interpolation='none')

plots[2,0].set_title('Line full width at half max')

colorbar(ab,title="[nm]")

plots[2,0].set_ylabel('Z [R$_\odot$]')

plots[2,0].set_xlabel('Y [R$_\odot$]')

############################################################

## remove unreasonable/outlier values from plotting range. Typical Fe XIII widths are 0.15. Outliers > 0.5 are excluded from scaling.

## PLOT INSIDE THE SAME SUBRANGES AS FULL WIDTHS above.

mx = np.ma.masked_array(specout[srx1:srx2,sry1:sry2,linen,9], mask=specout[srx1:srx2,sry1:sry2,linen,9] >= 0.50 )

vvmax=np.nanmax(mx)

ab=plots[2,1].imshow(specout[srx1:srx2,sry1:sry2,linen,9],extent=rnge,origin='lower',vmin=0.0,vmax=vvmax,cmap='YlGnBu',interpolation='none')

plots[2,1].set_title('Line non-thermal width')

colorbar(ab,title="[nm]")

#plots[2,1].set_ylabel('Z [R$_\odot$]')

plots[2,1].set_xlabel('Y [R$_\odot$]')

############################################################

## remove unreasonable/outlier values from plotting range. These should be a small addition to the total line widths. Change the 0.02 range as seems fit.

mx = np.ma.masked_array(specout[srx1:srx2,sry1:sry2,linen,10], mask=specout[srx1:srx2,sry1:sry2,linen,10] >= 0.020 )

vvmax=np.nanmax(mx)

ab=plots[2,2].imshow(specout[srx1:srx2,sry1:sry2,linen,10],extent=rnge,origin='lower',vmin=0.0,vmax=vvmax,cmap='YlOrRd',interpolation='none')

plots[2,2].set_title('Linear polarization fraction')

colorbar(ab,title="L / I ratio")

#plots[0,0].set_ylabel('Z [R$_\odot$]')

plots[2,2].set_xlabel('Y [R$_\odot$]')

############################################################

## remove unreasonable/outlier values from plotting range. These should be a small addition to the total line widths. Change the 0.02 range as seems fit.

mx = np.ma.masked_array(specout[srx1:srx2,sry1:sry2,linen,11], mask=specout[srx1:srx2,sry1:sry2,linen,11] >= 0.020 )

vvmax=np.nanmax(mx)

ab=plots[2,3].imshow(specout[srx1:srx2,sry1:sry2,linen,11],extent=rnge,origin='lower',vmin=0.0,vmax=vvmax,cmap='YlOrRd',interpolation='none')

plots[2,3].set_title('Total polarization fraction')

colorbar(ab,title="P / I ratio")

#plots[2,3].set_ylabel('Z [R$_\odot$]')

plots[2,3].set_xlabel('Y [R$_\odot$]')

############################################################

plt.suptitle(f"Selected line: {keyvals[4][linen]}")

plt.tight_layout()

### Save the putput plots

if not os.path.exists('./testrun_outputs'): ## make an output directory to keep things clean

os.makedirs('./testrun_outputs')

plt.savefig(f"./testrun_outputs/specout_1line_line{keyvals[4][linen]}_{datestamp}.pdf")

# ### plot 1-line BLOS outputs

# In[ ]:

fig, plots = plt.subplots(nrows=2, ncols=2, figsize=(12,6))

ab=plots[0,0].imshow(blosout[srx1:srx2,sry1:sry2,linen,0],extent=rnge,origin='lower',cmap='seismic',interpolation='none',vmin=-15.5,vmax=15.5)

plots[0,0].set_title('1$^{st}$ degenerate magnetograph solution')

colorbar(ab,title="B$_{LOS}$ [G]")

plots[0,0].set_ylabel('Z [R$_\odot$]')

#plots[0,0].set_xlabel('Y [R$_\odot$]')

ab=plots[0,1].imshow(blosout[srx1:srx2,sry1:sry2,linen,1],extent=rnge,origin='lower',cmap='seismic',interpolation='none',vmin=-15.5,vmax=15.5)

plots[0,1].set_title('2$^{nd}$ degenerate magnetograph solution')

colorbar(ab,title="B$_{LOS}$ [G]")

#plots[0,1].set_ylabel('Z [R$_\odot$]')

#plots[0,1].set_xlabel('Y [R$_\odot$]')

ab=plots[1,0].imshow(blosout[srx1:srx2,sry1:sry2,linen,2],extent=rnge,origin='lower',cmap='seismic',interpolation='none',vmin=-15.5,vmax=15.5)

plots[1,0].set_title('Classic magnetograph solution')

colorbar(ab,title="B$_{LOS}$ [G]")

plots[1,0].set_ylabel('Z [R$_\odot$]')

plots[1,0].set_xlabel('Y [R$_\odot$]')

ab=plots[1,1].imshow(blosout[srx1:srx2,sry1:sry2,linen,3],extent=rnge,origin='lower',cmap='twilight_shifted',interpolation='none',vmin=-np.pi/2,vmax=np.pi/2)

plots[1,1].set_title('$\Phi_B$ Magnetic Azimuth')

colorbar(ab,title="$\Phi_B$ [rad]")

plots[1,1].set_xlabel('Y [R$_\odot$]')

plt.suptitle(f"Selected line: {keyvals[4][linen]}")

plt.tight_layout()

### Save the putput plots

if not os.path.exists('./testrun_outputs'): ## make an output directory to keep things clean

os.makedirs('./testrun_outputs')

plt.savefig(f"./testrun_outputs/blosout_1line_line{keyvals[4][linen]}_{datestamp}.pdf")

# #### Check the issuemask for the selected pixel

# In[ ]:

xx=300

yy=500

# In[ ]:

importlib.reload(prepinv)

prepinv.iss_print(issuemask,tline=keyvals[4],xx=xx,yy=yy)

# #### Plot where a specific issue occurring over the entire map

# In[ ]:

importlib.reload(prepinv)

plt.figure(figsize=[8,8])

issue=32 ## Issue code to plot. see documentation for available codes:

issmap = prepinv.iss_print(issuemask,iss=issue,map1=True)

plt.imshow(issmap,extent=rnge,origin='lower',vmin=0.0,vmax=issue,cmap='hot_r',interpolation='none')

plt.ylabel('Z [R$_\odot$]')

plt.xlabel('Y [R$_\odot$]')

plt.tight_layout()