#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
# ATLAS FITTING

This is a CLI script to apply an atlas fitting to a given spectral image.

The purpose of this script is to select the lines in the atlas and data set used for dispersion fitting.
As a test for the result the script will present a plot with fitted continuum, fitted dispersion and the
remaining chi²-error. It will also create a txt file of the line positions used, in order to apply the
created config to a calibration set created by the spectroflat library.

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

## Execution
If the script is executable (preferred) you can run it directly as `./bin/setup --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/setup --help`.

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

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

import numpy as np
import yaml
from scipy.interpolate import CubicSpline
from spectroflat.smile import SmileInterpolator

sys.path.insert(0, os.path.join(sys.path[0], '..'))
from src.atlas_fit import Config
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, convolve_atlas_profile
from src.atlas_fit.plotting import *
from src.atlas_fit.fitting.line_finder import CommonLineFinder, LineTracker
from src.atlas_fit.utils import parse_shape, read_fits_data
from src.atlas_fit.fitting.wavelength import *
from src.atlas_fit.fitting.continuum import ContinuumCorrector, Detrender

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("-s", "--state", type=int, default=0, help='State to run the setup on')
aparser.add_argument("-r", "--refine", action="store_true", help='Refine selected lines (requires rangefile)')
args = aparser.parse_args()
log = logging.getLogger()

config = Config.from_yaml(args.path)

log.info('Loading spectrum (State=%s)...', args.state)
img = read_fits_data(config.input.corrected_frame)
print(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))


spectrum = img[args.state][roi].mean(axis=0)
if config.input.flipped:
    log.info('Flipping spectrum along wavelength axis ...')
    spectrum = np.flip(spectrum)
data = spectrum
if config.input.straighten:
    log.info('Temporarily remove strong gradients in spectrum ...')
    dt = Detrender(config, spectrum).run()
    plot_detrending(config, dt)
    data = dt.detrended

log.info('Loading atlas ...')
atlas = atlas_factory(config.atlas.key, config.input.wl_start, config.input.wl_end,
                      conversion=NM2ANGSTROM, path=config.atlas.path)
if config.input.model_psf is not None:
    atlas.intensity = convolve_atlas_with_model_PSF(atlas, config)
else:
    atlas.intensity = convolve_atlas_profile(atlas, config)

points = {"range": [], "lines_selected": [], "lines_deselected": []}
if args.range is not None:
    with open(args.range, 'r') as f:
        points = yaml.safe_load(f)
        # backwards compatibility & auto update
        if "lines" in points.keys():
            points["lines_selected"] = points["lines"]
            points["lines_deselected"] = []
            del(points["lines"])

if args.refine or args.range is None:
    points["range"] = AtlasRangeSelector(data, atlas, config).run(points["range"]).selection

if not points["range"]:
    log.warning('No points selected. Assuming MIN/MAX range')
    points = [atlas.wl.min(), atlas.wl.max()]
elif len(points["range"]) != 2:
    log.critical('Selected %s points, but exactly two are needed.', len(points))
    log.error('[ABORTED] Please start again and select the same two lines in both plots.')
    sys.exit(1)
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))

log.info('SELECTED ATLAS range:\n %.2f (%s) to %.2f (%s)', ap0_wl, ap0_px, ap1_wl, ap1_px)

atlas.clip(ap0_px, ap1_px)
clf = CommonLineFinder(data, atlas, config, points["lines_selected"], points["lines_deselected"]).run()
log.info('Common minima found: A=%s (D=%s)', len(clf.minima['atlas']), len(clf.minima['data']))
log.info('rejecting non-uniform lines (by bisector analysis)...')
clf.reject_non_uniform_lines()
log.info('Common usable minima found: A=%s (D=%s)', len(clf.minima['atlas']), len(clf.minima['data']))
ls = LineSelector(config, clf).run(args.refine)
if len(ls.selection) > 0 or len(ls.deselection) > 0:
    points["lines_selected"] = [float(atlas.wl_pos(px)) for px in ls.selection]
    points["lines_deselected"] = [float(atlas.wl_pos(px)) for px in ls.deselection]
    clf = ls.clf
    log.info("Force selected lines: %s", len(points["lines_selected"]))
    log.info("Force deselected lines: %s", len(points["lines_deselected"]))
    log.info('Updated number of minima: A=%s (D=%s)', len(clf.minima['atlas']), len(clf.minima['data']))

log.info('Refining line core positions based on Voigt profile fit.')
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('Refined core positions used: A=%s (D=%s)', len(refined_atlas_cores), len(refined_data_cores))

log.info('Finding wavelength solution...')
wl_anchors = atlas.wl_pos(refined_atlas_cores)
data_wl = wavelength_solution(refined_data_cores, wl_anchors, len(spectrum))

shifts = pixel_shifts_from_wavelength(data_wl)
measured_shifts = shifts
data_norm = spectrum / spectrum.mean()

log.info('Reloading atlas according to range wavelength solution: %.2f to %.2f', data_wl.min(), data_wl.max())
atlas = atlas_factory(config.atlas.key, data_wl.min(), data_wl.max(), conversion=NM2ANGSTROM, path=config.atlas.path)
atlas.intensity = convolve_atlas_with_model_PSF(atlas, config)

if config.debug:
    log.info('Plot wavelength solution result...')
    plot_wavelength_solution(config, atlas, spectrum, data_wl, shifts, refined_atlas_cores)

if not config.fitting.use_measured_shifts:
    deg = config.fitting.ideal_dispersion_degree
    shifts_fit = np.polynomial.Polynomial.fit(data_wl, shifts, deg)
    log.info('Using idealized shifts with deg=%s...', deg)
    shifts = shifts_fit(data_wl)

smax = np.abs(shifts.max())
dwl = (data_wl.max() - data_wl.min()) / len(data_wl)
log.info('Sampling: %.2e nm/px', dwl)
log.info('Apply wavelength correction to data (shift: %.1f px)', smax)
xes = np.arange(len(data_norm))
# Use '-shifts' as we want to UNDO the measured shifts
_, 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)

log.info('Determining continuum correction, gray stray light, and effective spectral FWHM...')
cc = ContinuumCorrector(data_lin, wl_lin, atlas, config).run()
print(repr(cc.result))

ref_wl = data_wl.mean()
log.info('Ref. WL: %.4e [nm]', ref_wl)
if config.input.model_psf is not None:
    log.info('Delta FWHM: %.4e nm (wrt to model PSF)', cc.delta_fwhm)
    log.info('Delta VINST: %.4e', cc.delta_fwhm * 10 * 299792.458 * np.sqrt(2 * np.log(2)) / (ref_wl * 20))
else:
    log.info('FWHM: %.4e', cc.data_fwhm)
    log.info('VINST: %.4e', cc.data_fwhm * 10 * 299792.458 * np.sqrt(2 * np.log(2)) / (ref_wl * 20))
log.info('STRAYLIGHT: %.3f [%%]  (gradient=%.1e, curvature=%.1e)',
         cc.gray_straylight, cc.gray_straylight_gradient, cc.gray_straylight_curve)
log.info('Plotting fitting result...')
plot_fit_result(config, cc, shifts)
plot_normalized_result(config, cc)

linefile = 'atlas_fine_range.yaml'
log.info('Writing used range to %s', linefile)
with open(linefile, 'w') as file:
    documents = yaml.dump(points, file)

if config.fitting.save_dispersion_function:
    dispersion_file = 'dispersion.txt'
    pix = np.arange(len(shifts), dtype=float)
    spx = np.polynomial.Polynomial.fit(pix, shifts, config.fitting.ideal_dispersion_degree)
    hdr = 'wl [px], wl [nm], measured dispersion [px], ideal dispersion [px]'
    data = np.array([pix, wl_lin, measured_shifts, spx(pix)])
    log.info('Writing dispersion to %s\n\tColumns: %s\n\tShape: %s', dispersion_file, hdr, data.shape)
    np.savetxt(dispersion_file, data, header=hdr)
