import numpy as np
import numpy.random as npr
import scipy.optimize as spo
import tomo_challenge.metrics as tcm

# custom data type, could be replaced with/tie in to tree.py class
# cut_vals is (nfeat, nbins - 1) numpy array, float
# tree_ids is ((nbins,) * nfeat) numpy array, int
TreePars = namedtuple('TreePars', ['cut_vals', 'tree_ids'])

# should maybe put this function in a class so we can call TreePars.to_array
def treepars_to_array(treepars):
    """
    Flattens cut_vals and tree_ids for optimizer
    """
    cuts = np.flatten(treepars.cut_vals)
    ids = np.flatten(treepars.tree_ids)
    arr = np.concatenate((cuts, ids))
    return(arr)

# should maybe put this function in a class so we can call TreePars.from_array
def array_to_treepars(arr):
    """
    Converts optimizer format of 1D array back into namedtuple of arrays
    """
    flat_cuts = arr[type(arr) == float]
    flat_ids = arr[type(arr) == int]
    nbins = len(np.unique(flat_ids))
    nfeat = len(flat_cuts) / (nbins - 1)
    # maybe do some assert checks with these just in case types have problems
    # cuts = arr[0:nfeat*(nbins-1)].reshape((nfeat, nbins-1))
    # ids = arr[feat*(nbins-1):].reshape((nbins,) * nfeat)
    cuts = flat_cuts.reshape((nfeat, nbins-1))
    ids = flat_ids.reshape((nbins,) * nfeat)
    treepars = TreePars(cuts, ids)
    return(treepars)

def get_cuts(galaxies, ival_treepars=None, nbins=3):
    """
    Obtains simplest possible bin definitions: cuts in the space of observables given number of bins

    Parameters
    ----------
    galaxies: numpy.ndarray, float
        observables (magnitudes and/or colors and/or errors) to serve as features for set of galaxies
        shape(galaxies) = (ngals, nfeat)
    ival_treepars: namedtuple, numpy.ndarray, float and int, optional
        initial values for decision tree parameters
        shape(ivals.cut_vals) = (nfeat, (nbins - 1))
        shape(tree_ids) = ((nbins,) * nfeat)
    nbins: int, optional
        number of bins for which to obtain cuts

    Returns
    -------
    assignments: numpy.ndarray, int
        bin assignment for each galaxy
        shape(assignments) = (ngals, 1)

    Notes
    -----
    `sort_gals` does the heavy lifting.
    `eval_metric` will call one of the metrics from [tomo_challenge](https://github.com/LSSTDESC/tomo_challenge/blob/master/tomo_challenge/metrics.py).
    The original idea for a general, non-cut-based optimizer was to have parameters equal to the (ngals) length array of ints representing the bin assignments, but that's not necessary for the simple cut-and-sweep barber and would probably break `spo.minimize`.
    """
    (ngals, nfeat) = np.shape(galaxies)

    if ival_treepars is None:
        cut_ivals = np.quantile(galaxies, np.linspace(0., 1., nbins), axis=1)
        assert(len(np.flatten(ivals)) == nbins**nfeat)

        # need structure and way of making dumb version of these
        tree_ids = npr.random_integers(0, nbins, nbins**nfeat)
        assert(len(np.unique(tree_ids)) == nbins)
        tree_ids.reshape((nfeat, nbins))

        ival_treepars = TreePars(cut_ivals, tree_ids)

    ivals = treepars_to_array(ival_treepars)
    opt_res = spo.minimize(eval_metric, ivals, args=galaxies)
    treepars = array_to_treepars(opt_res.x)

    assignments = sort_gals(galaxies, treepars)

    return(assignments)

def sort_gals(galaxies, tree_pars):
    """
    Divides available galaxies into subsets according to a given decision tree on their observables

    Parameters
    ----------
    galaxies: nfeature x n_gal array

    tree: tree object

    Notes
    -----
    could be based on bisect, or maybe a sklearn object?
    """
    pass


def eval_metric(arr, galaxies):
    """
    Just calls a metric from tomo_challenge wrapped for the `spo.minimize` API

    Notes
    -----
    Replace `tcm.metric` with actual call to one of the tomo_challenge metrics
    Actually, there's a problem in that the current tomo_challenge metrics require the true redshifts...
    """
    treepars = array_to_treepars(arr)
    assignments = sort_gals(galaxies, treepars)
    metval = tcm.metric(assignments)
    return metval

[((16, 11, 16, 40), 'numpy.flatten', 'np.flatten', ({(16, 22, 16, 39): 'treepars.cut_vals'}, {}), '(treepars.cut_vals)', True, 'import numpy as np\n'), ((17, 10, 17, 39), 'numpy.flatten', 'np.flatten', ({(17, 21, 17, 38): 'treepars.tree_ids'}, {}), '(treepars.tree_ids)', True, 'import numpy as np\n'), ((18, 10, 18, 37), 'numpy.concatenate', 'np.concatenate', ({(18, 25, 18, 36): '(cuts, ids)'}, {}), '((cuts, ids))', True, 'import numpy as np\n'), ((66, 21, 66, 39), 'numpy.shape', 'np.shape', ({(66, 30, 66, 38): 'galaxies'}, {}), '(galaxies)', True, 'import numpy as np\n'), ((80, 14, 80, 61), 'scipy.optimize.minimize', 'spo.minimize', (), '', True, 'import scipy.optimize as spo\n'), ((115, 13, 115, 36), 'tomo_challenge.metrics.metric', 'tcm.metric', ({(115, 24, 115, 35): 'assignments'}, {}), '(assignments)', True, 'import tomo_challenge.metrics as tcm\n'), ((28, 16, 28, 35), 'numpy.unique', 'np.unique', ({(28, 26, 28, 34): 'flat_ids'}, {}), '(flat_ids)', True, 'import numpy as np\n'), ((73, 19, 73, 62), 'numpy.random.random_integers', 'npr.random_integers', ({(73, 39, 73, 40): '0', (73, 42, 73, 47): 'nbins', (73, 49, 73, 61): 'nbins ** nfeat'}, {}), '(0, nbins, nbins ** nfeat)', True, 'import numpy.random as npr\n'), ((69, 42, 69, 68), 'numpy.linspace', 'np.linspace', ({(69, 54, 69, 56): '0.0', (69, 58, 69, 60): '1.0', (69, 62, 69, 67): 'nbins'}, {}), '(0.0, 1.0, nbins)', True, 'import numpy as np\n'), ((70, 19, 70, 36), 'numpy.flatten', 'np.flatten', ({(70, 30, 70, 35): 'ivals'}, {}), '(ivals)', True, 'import numpy as np\n'), ((74, 19, 74, 38), 'numpy.unique', 'np.unique', ({(74, 29, 74, 37): 'tree_ids'}, {}), '(tree_ids)', True, 'import numpy as np\n')]