from gempy.library import astromodels as am

import numpy as np
from matplotlib import pyplot as plt


def reduceScience(p):
    tweak_args = (p.user_params.pop("center", None),)
    p.prepare()
    p.addDQ()
    p.addVAR(read_noise=True)
    p.overscanCorrect()
    p.biasCorrect()
    p.ADUToElectrons()
    p.addVAR(poisson_noise=True)
    p.attachWavelengthSolution()
    p.flatCorrect()
    p.QECorrect()
    p.flagCosmicRays()
    p.distortionCorrect()
    p.findApertures()
    tweak_wavelength_solution(p.streams['main'], *tweak_args)
    p.skyCorrectFromSlit()
    p.adjustWCSToReference()
    p.resampleToCommonFrame(conserve=True)
    p.scaleCountsToReference()
    p.stackFrames()
    p.findApertures()
    p.traceApertures()
    p.storeProcessedScience(suffix="_2D")
    p.extractSpectra()
    p.fluxCalibrate()
    p.storeProcessedScience(suffix="_1D")


def tweak_wavelength_solution(adinputs, sky_row=None, r=1800):
    """
    Modify in-place the wavelength solution ("WAVE" model) of a mosaicked
    AstroDataGmos object with a user wavelength offset determined from a
    graphical comparison with a list of night-sky lines.

    Parameters
    ----------
    sky_row: int/None
        row to extract for sky spectrum (None => middle row)
    r: float
        resolving power to which to smooth the model sky spectrum
    """
    plt.ion()
    wsky, fsky = np.loadtxt("uves_sky.dat").T
    wsky *= 0.1
    init_shift = 0
    for ad in adinputs:
        if sky_row is None:
            sky_row = ad[0].shape[0] // 2
        else:
            sky_row = int(sky_row)
        data = np.ma.masked_array(ad[0].data[sky_row], ad[0].mask[sky_row])
        wave_model = am.get_named_submodel(ad[0].wcs.forward_transform, "WAVE")
        new_shift = 1

        while new_shift != 0:
            w = wave_model(np.arange(data.size)) + init_shift
            sky_model = np.sum([fs * np.exp(-0.5*((w-ws)/(ws/(r*2.35)))**2)
                                for ws, fs in zip(wsky, fsky)], axis=0)
            sky_model *= data.max() / sky_model.max()
            fig, ax = plt.subplots()
            ax.plot(w, data, 'k-')
            ax.plot(w, sky_model, 'b-')
            ax.set_title(ad.filename)
            plt.show()
            while True:
                new_shift = input("Enter wavelength shift (blue peak minus "
                                  "black peak): ")
                try:
                    new_shift = float(new_shift)
                except ValueError:
                    if new_shift == "":
                        new_shift = 0
                        break
                else:
                    break
            init_shift += new_shift
            plt.close(fig)

        wave_model.c0 += init_shift
        wave_model.inverse.domain = tuple(x + init_shift
                                          for x in wave_model.inverse.domain)