import argparse
import sys
import shlex
import numpy as np
import numpy.random
import scipy.spatial

from rvspecfit import make_interpol
from rvspecfit import serializer
import rvspecfit

git_rev = rvspecfit.__version__

INTERPOL_H5_NAME = 'interp_%s.h5'
INTERPOL_DAT_NAME = 'interpdat_%s.npy'


def getedgevertices(vec):
    """
    Given a set of n-dimentional points, return vertices of an n-dimensional
    cube that fully encompass/surrounds the data, This is sort of the envelope
    around the data

    Parameters
    ----------
    vec: numpy (Ndim, Npts)
        The array of input points

    Returns
    -------
    vec: numpy (Ndim, Nretpts)
        The returned array of surrounding points
    """

    pad = 0.2  # pad each dimension by this amount
    # (relative to the dimension width)
    ndim = len(vec[:, 0])
    span = np.ptp(vec, axis=1)
    lspans = vec.min(axis=1) - pad * span
    rspans = vec.max(axis=1) + pad * span
    # edges of the data
    positions = []
    # the number of vertices in the cube is 2**ndim
    for i in range(2**ndim):
        curpos = []
        for j in range(ndim):
            flag = (i & (2**j)) > 0
            curspan = [lspans[j], rspans[j]]
            curpos.append(curspan[flag])
        positions.append(curpos)
    positions = np.array(positions).T
    return positions


def execute(spec_setup,
            prefix=None,
            regular=False,
            perturb=True,
            revision='',
            cmdline=''):
    """
    Prepare the triangulation objects for the set of spectral data for a given
    spec_setup.

    Parameters
    ----------
    spec_setup: string
        The spectroscopic configuration
    prefix: string
        The prefix where the data are located and where the triangulation will
        be stored
    perturb: boolean
        Boolean flag whether to perturb a little bit the points before doing a
        triangulation. This prevents issues with degenerate vertices and
        stability of triangulation. Without perturbation find_simplex for
        example may revert to brute force search.

    Returns
    -------
    None

    """

    cur_fname = ('%s/' + make_interpol.SPECS_H5_NAME) % (prefix, spec_setup)
    D = serializer.load_dict_from_hdf5(cur_fname)
    (vec, specs, lam, parnames, mapper_module, mapper_class_name, mapper_args,
     lognorms, log_step) = (D['vec'], D['specs'], D['lam'], D['parnames'],
                            D['mapper_module'], D['mapper_class_name'],
                            D['mapper_args'], D['lognorms'], D['log_step'])
    del D
    mapper = make_interpol.get_mapper(mapper_module, mapper_class_name,
                                      mapper_args)

    vec = vec.astype(float)
    vec = mapper.forward(vec)
    if not np.isfinite(vec.sum()):
        raise RuntimeError('Something is broken the parameters are not finite')
    ndim = len(vec[:, 0])
    ret_dict = {}

    if not regular:
        perturbation_amplitude = 1e-6
        # It turn's out that Delaunay is sometimes unstable when dealing with
        # regular grids, so perturb points
        if perturb:
            state = np.random.get_state()
            # Make it deterministic
            np.random.seed(1)
            vec = vec + np.random.uniform(-perturbation_amplitude,
                                          perturbation_amplitude,
                                          size=vec.shape)
            np.random.set_state(state)

        # get the positions that are outside the existing grid
        edgepositions = getedgevertices(vec)
        nedgepos = len(edgepositions.T)

        # find out the nearest neighbors for those outside points
        nearnei = scipy.spatial.cKDTree(vec.T).query(edgepositions.T)[1]
        vec = np.hstack((vec, edgepositions))

        nspec, lenspec = specs.shape

        # add nearest neighbor sectra to the grid at the edge locations
        specs = np.append(specs, np.array([specs[_] for _ in nearnei]), axis=0)

        # extra flags that allow us to detect out of the grid cases (i.e inside
        # our grid the flag should be 0)
        extraflags = np.concatenate((np.zeros(nspec), np.ones(nedgepos)))

        lognorms = np.concatenate((lognorms, np.zeros(nedgepos)))
        vec = vec.astype(np.float64)
        extraflags = extraflags.astype(np.float64)
        specs = specs.astype(np.float64)
        extraflags = extraflags[:, None]

        triang = scipy.spatial.Delaunay(vec.T)
        ret_dict['triang'] = triang
        ret_dict['extraflags'] = extraflags
        ret_dict['interpolation_type'] = 'triangulation'
    else:
        uvecs0 = [
            np.unique(vec[i, :], return_inverse=True) for i in range(ndim)
        ]
        # Unique grid positions for each dimension
        uvecs = [_[0] for _ in uvecs0]
        # location inside uvec
        vecids = [_[1] for _ in uvecs0]
        # lens is the list of sizes [10,20, 30] with the length
        # being the number of dims
        lens = [len(_) for _ in uvecs]
        # these will be the locations in the input arrays on the grid
        # ie if we provide 3 points for [0,0], [0,1], [1,1]
        # the idgrid will be [[0,1],[-1,2]]
        idgrid = np.zeros(lens, dtype=int) - 1
        idgrid[tuple(vecids)] = np.arange(vec.shape[1])
        ret_dict['uvecs'] = uvecs
        ret_dict['regular'] = True
        ret_dict['idgrid'] = idgrid
        ret_dict['interpolation_type'] = 'regulargrid'
    savefile = ('%s/' + INTERPOL_H5_NAME) % (prefix, spec_setup)
    ret_dict['lam'] = lam
    ret_dict['log_step'] = log_step
    ret_dict['vec'] = vec
    ret_dict['parnames'] = parnames
    ret_dict['mapper_class_name'] = mapper_class_name
    ret_dict['mapper_module'] = mapper_module
    ret_dict['mapper_args'] = mapper_args
    ret_dict['revision'] = revision
    ret_dict['lognorms'] = lognorms
    ret_dict['git_rev'] = git_rev
    ret_dict['cmdline'] = cmdline

    # for triangular/grid I'm allowing pickling
    serializer.save_dict_to_hdf5(savefile, ret_dict, allow_pickle=True)
    np.save(('%s/' + INTERPOL_DAT_NAME) % (prefix, spec_setup),
            np.ascontiguousarray(specs))


def main(args=None):
    if args is None:
        args = sys.argv[1:]
    cmdline = shlex.join([sys.argv[0]] + list(args))
    parser = argparse.ArgumentParser(
        description='Create N-D spectral interpolation files')
    parser.add_argument(
        '--prefix',
        type=str,
        help='Location of the interpolated and convolved input spectra',
        required=True)
    parser.add_argument('--setup',
                        type=str,
                        help='Name of the spectral configuration',
                        required=True)
    parser.add_argument('--regulargrid',
                        action='store_true',
                        help='Use regular grid interpolation ',
                        required=False)
    parser.add_argument('--revision',
                        type=str,
                        help='Revision of the data files/run',
                        required=False)

    args = parser.parse_args(args)
    execute(args.setup,
            prefix=args.prefix,
            revision=args.revision,
            regular=args.regulargrid,
            cmdline=cmdline)


if __name__ == '__main__':
    main()
