#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
# SPECTROFLAT CALIBRATION DATA AMENDMENT

This is a CLI script to amend the calibration data from the `spectroflat` library.
It will add dispersion correction (and chi^2 error) to the `OffsetMap`.

## Configuration
The configuration is given as a YAML file. See in examples folder or `atlas_fit/config.py`
for documentation on parameters.

## Execution
If the script is executable (preferred) you can run it directly as `./bin/amend_spectroflat --help`.
The script should automatically be marked as executable after git checkout.
If this is not the case and/or you do not want this, you can execute it with
your favorite python interpreter as `python3 bin/amend_spectroflat --help`.

See `--help` switch for further usage details.

## Author(s)
hoelken@mps.mpg.de
"""
import logging
import os
import sys
from argparse import ArgumentParser

import numpy as np
import yaml
from astropy.io import fits
from scipy.interpolate import pchip_interpolate
from scipy.ndimage import gaussian_filter, zoom
from spectroflat.smile import OffsetMap, SmileInterpolator
from spectroflat.utils.processing import MP
from qollib.ui.progress import bar

sys.path.insert(0, os.path.join(sys.path[0], '..'))
from src.atlas_fit.base.constants import NM2ANGSTROM
from src.atlas_fit.atlantes import atlas_factory
from src.atlas_fit.models.differential_convolution import convolve_atlas_with_model_PSF
from src.atlas_fit.fitting.line_finder import CommonLineFinder, LineTracker
from src.atlas_fit.fitting.wavelength import *
from src.atlas_fit.fitting.continuum import ContinuumCorrector, Detrender
from src.atlas_fit import Config
from src.atlas_fit.utils import parse_shape, read_fits_data, compute_continuum_correction

aparser = ArgumentParser(description='Fit a Solar Atlas to the data')
aparser.add_argument('path', type=str, metavar='CONFIG_PATH',
                     help='Path to the configuration file to use')
aparser.add_argument('range', nargs='?', metavar='RANGE_PATH', default=None,
                     help='Path to file with list of lines to use (optional)')
aparser.add_argument("-o", "--out", type=str, nargs="?", default=".", help='Output folder. default: current folder')
args = aparser.parse_args()
log = logging.getLogger()

config = Config.from_yaml(args.path)
with open(args.range, 'r') as f:
    points = yaml.safe_load(f)
    # backwards compatibility
    if "lines" in points.keys():
        points["lines_selected"] = points["lines"]
        points["lines_deselected"] = []

log.info('Loading atlas ...')
atlas = atlas_factory(config.atlas.key, config.input.wl_start, config.input.wl_end,
                      conversion=NM2ANGSTROM, path=config.atlas.path)
atlas.intensity = convolve_atlas_with_model_PSF(atlas, config)

ap0_wl, ap1_wl = points["range"]
ap0_px = np.argmin(np.abs(atlas.wl - ap0_wl))
ap1_px = np.argmin(np.abs(atlas.wl - ap1_wl))
atlas.clip(ap0_px, ap1_px)
log.info('SELECTED ATLAS range:\n %.2f (%s) to %.2f (%s)', ap0_wl, ap0_px, ap1_wl, ap1_px)

log.info('Loading SPECTROFLAT calibration data.')
img = read_fits_data(config.input.corrected_frame)
offsets = OffsetMap.from_file(config.input.offset_map)
if config.input.flipped:
    offsets.map = -1 * np.flip(offsets.map, axis=2)
    img = np.flip(img, axis=2)
    roi = parse_shape(config.input.correction_roi)
    if None not in [roi[1].start, roi[1].stop] and ':' not in str(roi[1]):
        config.input.correction_roi = f'[{roi[0]}, {img.shape[2] - roi[1].stop}:{img.shape[2] - roi[1].start}]'
    print(config.input.correction_roi)

log.info('Amending OffsetMap %s', config.input.offset_map)
log.info('=== Fitting dispersion and continuum ===')
delta_offsets = []
orig_means = []
correction_state = []
wl_vectors = []
straylight = []
straylight_functions = []
fwhm = []

states, rows, cols = img.shape
roi = parse_shape(config.input.correction_roi) if config.input.correction_roi is not None else None
if config.input.flipped and roi is not None:
    stop = img.shape[2] if roi[1].stop is None else roi[1].stop
    start = 0 if roi[1].start is None else roi[1].start
    roi = (roi[0], slice(img.shape[2] - stop, img.shape[2] - start))
if roi is None:
    roi = (slice(None, None), slice(None, None))
start = roi[1].start if roi[1].start is not None else 0
stop = roi[1].stop if roi[1].stop is not None else cols
r0 = roi[0].start if roi[0].start is not None else 0
r1 = roi[0].stop if roi[0].stop is not None else rows
rows = np.arange(rows)
log.info("Correction ROI: %s", repr(roi))

for s in range(config.input.mod_states):
    log.info(f'::: State {s} of {config.input.mod_states - 1} (ROI=%s)', roi)
    spectrum = img[s][roi].mean(axis=0)
    state_mean = spectrum.mean()
    if config.input.straighten:
        data = Detrender(config, spectrum).run().detrended
    else:
        data = spectrum

    clf = CommonLineFinder(data, atlas, config, points["lines_selected"], points["lines_deselected"]).run()
    clf.reject_non_uniform_lines()
    lt = LineTracker(config, spectrum, atlas, clf.minima).fit_exact_positions(points["lines_selected"])
    refined_atlas_cores = lt.refined_atlas_cores
    refined_data_cores = lt.refined_data_cores
    log.info(f"\tCommon minima found: {len(refined_atlas_cores)}")
    wl_anchors = atlas.wl_pos(refined_atlas_cores)
    data_wl = wavelength_solution(refined_data_cores, wl_anchors, len(spectrum))
    # Use '-shifts' as we want to UNDO the measured shifts
    shifts = -pixel_shifts_from_wavelength(data_wl)
    log.info(f"\tmax shift= {shifts[np.argmax(np.abs(shifts))]:.3f} px,")

    state_atlas = atlas_factory(config.atlas.key, data_wl.min(), data_wl.max(),
                                conversion=NM2ANGSTROM, path=config.atlas.path)
    state_atlas.intensity = convolve_atlas_with_model_PSF(state_atlas, config)
    if not config.fitting.use_measured_shifts:
        shifts_fit = np.polynomial.Polynomial.fit(data_wl, shifts, config.fitting.ideal_dispersion_degree)
        shifts = shifts_fit(data_wl)

    data_norm = spectrum / spectrum.mean()
    xes = np.arange(len(data_norm))
    _, data_lin = SmileInterpolator.desmile_row((0, xes, shifts, data_norm))

    wl_lin = anchored_linear_wavelength_grid(data_wl, data_lin, atlas, wl_anchors, -shifts, refined_data_cores)
    cc = ContinuumCorrector(data_lin, wl_lin, state_atlas, config).run()
    log.info(f"\tFWHM={cc.data_fwhm:.2e}nm")
    log.info(f"\tstray={cc.gray_straylight:.1f}% ({state_mean*(cc.gray_straylight/100):.1f}cts of {state_mean:.1f}cts)")
    delta_offsets.append(shifts)
    wl_vectors.append(wl_lin)
    orig_means.append(state_mean)
    fwhm.append(cc.data_fwhm)
    straylight.append(cc.gray_straylight)
    straylight_functions.append(cc.straylight_function)

    log.info(f"\tComputing continuum correction for all rows from {r0} to {r1}")
    cont_corr = np.ones_like(img[s])
    arguments = [(r, r1, img[s, r, start:stop], cc, xes, shifts, state_mean) for r in range(r0, r1, 1)]
    res = dict(MP.simultaneous(compute_continuum_correction, arguments))
    cont_corr[r0:r1, start:stop] = np.array([res[row] for row in range(r0, r1, 1)])
    correction_state.append(cont_corr)

log.info('Post processing results...')
if config.fitting.squash_offsets:
    log.info('Averaging modulated offsets to a global state (squash).')
    offsets.squash()
    wl = np.average(wl_vectors, axis=0)
    delta_offsets = np.repeat(np.array(delta_offsets).mean(axis=0, keepdims=True), config.input.mod_states, axis=0)
else:
    log.warning('Providing MODULATED OFFSETS: Check demodulated data for artifacts!')
    wl_vectors = np.array(wl_vectors)
    delta_offsets = np.array(delta_offsets)
    max_state = np.argmax(wl_vectors[:, 0])
    min_state = np.argmax(wl_vectors[:, 0])
    log.info('Post aligning WL. Max delta: %.3e', wl_vectors[max_state, 0] - wl_vectors[min_state, 0])
    wl = wl_vectors[max_state]
    max_start = wl_vectors[max_state, 0]
    for s in range(config.input.mod_states):
        dispersion = (wl_vectors[s, 1:11] - wl_vectors[s, 0:10]).mean()
        delta_wl = max_start - wl_vectors[s, 0]
        delta_offsets[s] = delta_offsets[s] + delta_wl * dispersion
        log.info('State %s WL shift correction: %.3e', s, delta_wl * dispersion)

log.info('=== Add WL calibration to offset map ===')
out_file = os.path.join(args.out, 'wl_calibrated_offsets.fits')
log.info('Creating "%s"', out_file)

log.debug('Input correction roi: %s', config.input.correction_roi)
for s in range(config.input.mod_states):
    for r in range(r0, r1, 1):
        offsets.map[s, r, start:stop] = offsets.map[s, r, start:stop] + delta_offsets[s]

rpx = np.abs(stop - start) // 2
WCS = {
    'BUNIT': ('', 'Unit of original pixel value'),
    'BTYPE': ('Relative Intensity', 'Type of  pixel value'),
    'WAVELEN': (wl.mean(), 'Center wavelength (nm)'),
    'WAV-MAX': (max(wl), 'Maximum wavelength recorded (nm)'),
    'WAV-MIN': (min(wl), 'Minimum wavelength recorded (nm)'),
    'C2YPE3': ('WAVELENGTH', 'Type of projection used for #3 axis in 2nd WCS'),
    'C2PIX3': (rpx, 'Reference pixel in axis #3'),
    'C2VAL3': (wl[rpx - min(start, stop)], 'Physical value of ref pix in axis #3 for WCS'),
    'C2ELT3': ((wl[-1] - wl[0]) / np.abs(stop - start), 'Size projected into a detector pixel'),
}

offsets.header['HIERARCH MIN_WL_NM'] = min(wl)
offsets.header['HIERARCH MAX_WL_NM'] = max(wl)
offsets.header['HIERARCH MIN_WL_PX'] = stop if config.input.flipped else start
offsets.header['HIERARCH MAX_WL_PX'] = start if config.input.flipped else stop
offsets.header['HIERARCH DISPERSION'] = (wl[-1] - wl[0]) / np.abs(stop - start)
offsets.header['HIERARCH STRAYLIGHT'] = np.mean(straylight)
offsets.header['HIERARCH FWHM'] = np.mean(fwhm)
for k, v in WCS.items():
    offsets.header[k] = v[0]
log.info("WAVELENGTH INFO:\n%s", repr(offsets.header))
offsets.wl = wl
if config.input.flipped:
    offsets.header.append(('WL_FLIPPED', True), 'Original map was flipped in wavelength direction')
    axis = 1 if offsets.is_squashed() else 2
    offsets.map = -1 * np.flip(offsets.map, axis=axis)
offsets.dump(out_file=out_file)

log.info('=== Add continuum calibration to illumination pattern ===')
out_file = os.path.join(args.out, 'amended_soft_flat.fits')
log.info('Creating "%s"', out_file)

log.info('Compute correction image...')
correction_state = np.array(correction_state)
illum_pattern = fits.getdata(config.input.soft_flat)
if config.input.flipped:
    illum_pattern = np.flip(illum_pattern, axis=2)

xes = np.arange(stop - start)
log.info('Applying additional offset correction to original illumination pattern (HDU=0)')
for s in range(states):
    arguments = [(r, xes, delta_offsets[s], illum_pattern[s, r, start:stop]) for r in rows]
    res = dict(MP.simultaneous(SmileInterpolator.desmile_row, arguments))
    illum_pattern[s, :, roi[1].start:roi[1].stop] = np.array([res[row] for row in rows])
if config.input.flipped:
    illum_pattern = np.flip(illum_pattern, axis=2)

illum_pattern = fits.PrimaryHDU(illum_pattern, fits.getheader(config.input.soft_flat))
for k, v in WCS.items():
    illum_pattern.header.append(k, v[0], v[1])

log.info('Adding 2D continuum correction surface (HDU=1)')

if config.input.flipped:
    correction_state = np.flip(correction_state, axis=2)
cont = fits.ImageHDU(correction_state)
aim = atlas.intensity.mean()
cont.header.append(('HIERARCH MEAN_ATLAS_CONTINUUM', f"{atlas.continuum.mean():.5e}", "W cm-1 ster-1 A-1"))
cont.header.append(('HIERARCH MEAN_ATLAS_INTENSITY', f"{aim:.5e}", "W cm-1 ster-1 A-1"))
for s in range(states):
    cont.header.append((f'HIERARCH MEAN_INTENSITY_CONVERSION_{s}', f"{aim / orig_means[s]}"))
    cont.header.append((f'HIERARCH FF_MEAN_INTENSITY_{s}', f"{orig_means[s]}"))

log.info('Adding 2D Straylight surface (HDU=2)')
stray_pat = np.zeros((illum_pattern.data.shape[1], illum_pattern.data.shape[2]))
stray_pat[r0:r1, start:stop] = np.array(straylight_functions).mean(axis=0)
stray_perc = straylight
straylight = fits.ImageHDU(stray_pat)
straylight.header['HIERARCH MEAN STRAYLIGHT'] = offsets.header['HIERARCH STRAYLIGHT']
for i in range(len(stray_perc)):
    straylight.header[f"HIERARCH STATE_{i}_STRAYLIGHT"] = stray_perc[i]
    straylight.header[f"HIERARCH STATE_{i}_STRAYLIGHT_CTS"] = stray_perc[i]/100 * orig_means[i]
print(repr(straylight.header))

log.info('Adding 1D wavelength information (HDU=3)')
wlhdu = fits.ImageHDU(wl)
for k, v in WCS.items():
    wlhdu.header[k] = v[0]
wlhdu.header.append(('WLSAMP', f'{(max(wl) - min(wl)) / len(wl):.8e}', 'sampling'))
wlhdu.header.append(('NWL', len(wl), 'WL-points'))

if config.input.flipped:
    illum_pattern.header.append(('WL_FLIPPED', True, 'Original map was sp flipped'))
    cont.header.append(('WL_FLIPPED', True, 'Original map was sp flipped'))
    straylight.header.append(('WL_FLIPPED', True, 'Original map was sp flipped'))
    wlhdu.header.append(('WL_FLIPPED', True, 'Original map was sp flipped'))

hdul = fits.HDUList([illum_pattern, cont, straylight, wlhdu])
hdul.writeto(out_file, overwrite=True)

log.info('=== Applying calibration to flat-fielded flat ===')
out_file = os.path.join(args.out, 'amended_corrected_frame.fits')
log.info('Creating "%s"', out_file)

corr_frame = fits.getdata(config.input.corrected_frame)
if config.input.flipped:
    corr_frame = np.flip(corr_frame, axis=2)

log.info('\t- (1) Applying additional offset correction...')
for s in range(states):
    arguments = [(r, xes, delta_offsets[s], corr_frame[s, r, start:stop]) for r in rows]
    res = dict(MP.simultaneous(SmileInterpolator.desmile_row, arguments))
    corr_frame[s, :, roi[1].start:roi[1].stop] = np.array([res[row] for row in rows])
    stray = straylight.data.copy()
    for r in range(r0, r1, 1):
        stray[r] *= corr_frame[s, r, start:stop].mean()
    corr_frame[s] = (corr_frame[s] - stray) / (cont.data[s] / np.mean(cont.data))  # we stay in [DN]
if config.input.flipped:
    corr_frame = np.flip(corr_frame, axis=2)

log.info('\t- (2) Applying continuum & straylight correction...')
cal_corr_frame = fits.PrimaryHDU(corr_frame, fits.getheader(config.input.corrected_frame))

log.info('\t- (3) Modifying header...')
cal_corr_frame.header['HIERARCH MIN_WL_NM'] = offsets.header['HIERARCH MIN_WL_NM']
cal_corr_frame.header['HIERARCH MIN_WL_PX'] = offsets.header['HIERARCH MIN_WL_PX']
cal_corr_frame.header['HIERARCH MAX_WL_NM'] = offsets.header['HIERARCH MAX_WL_NM']
cal_corr_frame.header['HIERARCH MAX_WL_PX'] = offsets.header['HIERARCH MAX_WL_PX']
cal_corr_frame.header['HIERARCH DISPERSION'] = offsets.header['HIERARCH DISPERSION']
cal_corr_frame.header['HIERARCH WLSAMP'] = wlhdu.header['WLSAMP']
cal_corr_frame.header['HIERARCH NWL'] = wlhdu.header['NWL']
cal_corr_frame.header['HIERARCH MEAN_ATLAS_CONTINUUM'] = cont.header['HIERARCH MEAN_ATLAS_CONTINUUM']
cal_corr_frame.header['HIERARCH MEAN_ATLAS_INTENSITY'] = cont.header['HIERARCH MEAN_ATLAS_INTENSITY']
for s in range(states):
    cal_corr_frame.header[f'HIERARCH MEAN_INTENSITY_CONVERSION_{s}'] = cont.header[
        f'HIERARCH MEAN_INTENSITY_CONVERSION_{s}'
    ]
if config.input.flipped:
    cal_corr_frame.header.append(('WL_FLIPPED', True), 'Original map was sp flipped')

hdul = fits.HDUList([cal_corr_frame, wlhdu])
hdul.writeto(out_file, overwrite=True)

log.info('=== Generating generic atlas reference file from corrected frame ===')
out_file = os.path.join(args.out, 'amended_corrected_frame_average.txt')
log.info('Creating "%s"', out_file)
print(corr_frame.shape)
data = corr_frame[:, roi[0].start:roi[0].stop, roi[1].start:roi[1].stop].mean(axis=(0, 1))
cont = np.ones(data.shape) * np.max(data) + 1e-3
print(corr_frame.shape)
wl = wlhdu.data * 10  # in \AA
np.savetxt(out_file, np.array([wl, data, cont]).T)

log.info('Wavelength and continuum calibration: Done.')
