import numpy as np
import scipy.optimize
import scipy.interpolate
from rvspecfit import make_ccf
from rvspecfit.spec_fit import SpecData
from rvspecfit import serializer
import logging

# Number of templates processed together in the batched FFT path.
_CCF_FFT_BATCH_SIZE = 256


class CCFCache:
    """ Singleton caching CCF information """
    ccf_info = {}
    ccfs = {}
    ccf2s = {}
    ccf_models = {}


def get_ccf_info(spec_setup, config):
    """
    Returns the CCF info from the pickled file for a given spectroscopic
    setup

    Parameters
    ----------
    spec_setup: string
        The spectroscopic setup needed
    config: dict
        The dictionary with the config

    Returns
    -------
    d: dict
        The dictionary with the CCF Information as saved by the make_ccf code

    """
    if spec_setup not in CCFCache.ccfs:
        prefix = config['template_lib'] + '/'
        ccf_continuum = config.get('ccf_continuum_normalize')
        if ccf_continuum is None:
            ccf_continuum = True
        ccf_info_fname = prefix + make_ccf.get_ccf_info_name(
            spec_setup, ccf_continuum)
        ccf_dat_fname = prefix + make_ccf.get_ccf_dat_name(
            spec_setup, ccf_continuum)
        ccf_mod_fname = prefix + make_ccf.get_ccf_mod_name(
            spec_setup, ccf_continuum)

        CCFCache.ccf_info[spec_setup] = serializer.load_dict_from_hdf5(
            ccf_info_fname)
        C = np.load(ccf_dat_fname, mmap_mode='r')
        CCFCache.ccfs[spec_setup] = C['fft']
        CCFCache.ccf2s[spec_setup] = C['fft2']
        CCFCache.ccf_models[spec_setup] = np.load(ccf_mod_fname, mmap_mode='r')
    return CCFCache.ccfs[spec_setup], CCFCache.ccf2s[
        spec_setup], CCFCache.ccf_models[spec_setup], CCFCache.ccf_info[
            spec_setup]


def fit(specdata, config):
    """
    Process the data by doing cross-correlation with templates

    Parameters
    ----------
    specdata: list of SpecData objects
        The list of data that needs to be fitted from different spectral
        setups.
    config: dict
        The configuration dictionary

    Returns
    -------
    results: dict
        The dictionary with results such as best template parameters,
        best velocity, best vsini.

    """
    # configuration parameters

    maxvel = config.get('max_vel') or 1000
    # only search for CCF peaks from -maxvel to maxvel
    nvelgrid = 2 * int(maxvel * 1. / (config.get('vel_step0') or 2)) + 1
    # number of points on the ccf in the specified velocity range
    vel_grid = np.linspace(-maxvel, maxvel, nvelgrid)

    # Algorithm outline:
    # 1) Preprocess each input arm once and cache FFT products needed for CCF.
    # 2) Loop over template chunks, compute CCFs in batch, interpolate each arm
    #    to the shared velocity grid, and sum chi-square contributions.
    velstep = {}  # step in the ccf in velocity
    ccf_dats = {}  # ffts of templates
    ccf2_dats = {}  # ffts of templates of squared spectra
    ccf_infos = {}  # ccf configurations
    ccf_mods = {}  # the actual template models
    proc_specs = {}  # actual data processed/continuum normalized etc
    proc_ivars = {}
    setup_states = {}
    setups = []
    total_sse = 0
    if isinstance(specdata, SpecData):
        # if we got a single one put it in the list
        specdata = [specdata]
    for cursd in specdata:
        spec_setup = cursd.name
        setups.append(spec_setup)
        lam = cursd.lam
        spec = cursd.spec
        espec = cursd.espec
        (ccf_dats[spec_setup], ccf2_dats[spec_setup], ccf_mods[spec_setup],
         ccf_infos[spec_setup]) = get_ccf_info(spec_setup, config)
        ccfconf = ccf_infos[spec_setup]['ccfconf']
        logl0 = ccfconf['logl0']
        logl1 = ccfconf['logl1']
        npoints = ccfconf['npoints']
        proc_spec, proc_ivar = make_ccf.preprocess_data(lam,
                                                        spec,
                                                        espec,
                                                        badmask=cursd.badmask,
                                                        ccfconf=ccfconf)
        proc_specs[spec_setup] = proc_spec
        proc_ivars[spec_setup] = proc_ivar
        total_sse += (proc_spec**2 * proc_ivar).sum()
        spec_fft = np.fft.rfft(proc_spec * proc_ivar)
        ivar_fft = np.fft.rfft(proc_ivar)

        spec_fftconj = spec_fft.conj()
        ivar_fftconj = ivar_fft.conj()

        cur_step = (np.exp((logl1 - logl0) / npoints) - 1) * 3e5
        lspec = len(proc_spec)
        # Importantly this has to be length of the spectrum rather than
        # length of the fft since we are using rfft
        cur_off = lspec // 2
        # this is the wrapping point
        cur_vels = -((np.arange(lspec) + cur_off) % lspec - cur_off) * cur_step
        # now cur_vels[lspec-off] is the first positive velocity
        # we need to np.roll(X,cur_off)  to make it continuous
        # notice that it is decreasing and it corresponds to the velocity of
        # the ccf pixels
        cur_ind = (np.abs(cur_vels) < (maxvel + cur_step))
        # boolean mask within the required velocity range
        assert (cur_ind.sum() % 2 == 1)  # must be odd
        cur_ind = np.roll(np.nonzero(cur_ind)[0], cur_ind.sum() // 2)
        # these are indices that makes it monotonic
        cur_ind = cur_ind[::-1]
        # that provides indices that will go from negative
        # to positive velocities
        velstep[spec_setup] = cur_step
        cur_sub_vels = cur_vels[cur_ind]
        if not np.all(np.diff(cur_sub_vels) > 0):
            raise RuntimeError('Velocity grid for CCF interpolation is invalid')
        setup_states[spec_setup] = dict(spec_fftconj=spec_fftconj,
                                        ivar_fftconj=ivar_fftconj,
                                        ccf_dats=ccf_dats[spec_setup],
                                        ccf2_dats=ccf2_dats[spec_setup],
                                        subind=cur_ind,
                                        vels=cur_sub_vels,
                                        continuum=ccfconf['continuum'])

    ref_setup = setups[0]
    ref_info = ccf_infos[ref_setup]
    for spec_setup in setups[1:]:
        cur_info = ccf_infos[spec_setup]
        mismatch = (
            ref_info['parnames'] != cur_info['parnames']
            or not np.array_equal(ref_info['params'], cur_info['params'])
            or not np.array_equal(ref_info['vsinis'], cur_info['vsinis']))
        if mismatch:
            raise RuntimeError('The parameters of the CCF templates do not match')

    # the logic is the following
    # if array y is shifted by n pixels to the right side wrt x
    # ifft(fft(x)*fft(y).conj) will peak at pixel N-n (0based)
    # or if array is shifted to n pixels to the left it will peak at n (0based)

    nffts = [setup_states[setup]['ccf_dats'].shape[0] for setup in setups]
    if len(set(nffts)) != 1:
        raise RuntimeError('CCF template counts are inconsistent across setups')
    nfft = nffts[0]
    nvel = len(vel_grid)
    all_chisqs = np.zeros((nfft, nvel))

    # We evaluate templates in chunks to reduce Python overhead while keeping
    # memory use bounded. For each arm we compute CCFs for a whole chunk, then
    # interpolate them onto the common velocity grid and accumulate across arms.
    for start in range(0, nfft, _CCF_FFT_BATCH_SIZE):
        stop = min(start + _CCF_FFT_BATCH_SIZE, nfft)
        slc = slice(start, stop)
        cur_chisq = np.zeros((stop - start, nvel))
        for spec_setup in setups:
            state = setup_states[spec_setup]
            curf = state['ccf_dats'][slc, :]
            curf2 = state['ccf2_dats'][slc, :]
            curccf0 = np.fft.irfft(curf * state['spec_fftconj'][None, :], axis=1)
            curccf1 = np.fft.irfft(curf2 * state['ivar_fftconj'][None, :], axis=1)
            # chisquare i -2* l * S/E^2 xx T + l^2 1/E^2 xx T^2
            #  l is the multiplier (S-lT)
            # the best value is l = ((S/E^2) xx T)/(1/E^2 xx T^2)
            # Thus  the best chisq is -((S/E^2) xx T)^2/(1/E^2 xx T^2)
            # where xx is the convolution operator
            if state['continuum']:
                chisq = -2 * curccf0 + curccf1
            else:
                chisq = (-curccf0**2 / curccf1)
            cur_chisq += scipy.interpolate.interp1d(state['vels'],
                                                    chisq[:, state['subind']],
                                                    kind='linear',
                                                    axis=1,
                                                    assume_sorted=True)(
                                                        vel_grid)
            # we interpolate all the ccf from every arm
            # to the same velocity grid
        all_chisqs[slc, :] = cur_chisq

    all_chisqs += total_sse
    best_id = np.argmin(all_chisqs.min(axis=1))
    best_ccf = all_chisqs[best_id]
    best_pix = np.argmin(best_ccf)
    if best_pix not in [0, len(best_ccf) - 1]:
        coeffs = np.polyfit(vel_grid[best_pix - 1:best_pix + 2],
                            best_ccf[best_pix - 1:best_pix + 2],
                            deg=2)
        if coeffs[0] > 0:
            best_vel = -coeffs[1] / (2 * coeffs[0])
        else:
            best_vel = vel_grid[best_pix]
        # -b/2a
    else:
        best_vel = vel_grid[best_pix]

    if not np.isfinite(all_chisqs[best_id, best_pix]):
        logging.error('Cross-correlation failed')
        raise RuntimeError('Cross-correlation step failed')

    best_model = {}
    for spec_setup in setups:
        best_model[spec_setup] = np.roll(ccf_mods[spec_setup][best_id],
                                         int(best_vel / velstep[spec_setup]))
    best_par = ccf_infos[setups[0]]['params'][best_id]
    best_par = dict(zip(ccf_infos[setups[0]]['parnames'], best_par))
    best_vsini = ccf_infos[setups[0]]['vsinis'][best_id]

    result = dict(best_par=best_par,
                  best_vel=best_vel,
                  best_ccf=best_ccf,
                  best_vsini=best_vsini,
                  best_model=best_model,
                  proc_spec=proc_specs,
                  vel_grid=vel_grid)
    return result
