#
# Auxiliary functions for visualizations.
#
from __future__ import absolute_import, division
from __future__ import print_function, unicode_literals
import numpy as np


def generate_grid(lower, upper, splits, fixed=[]):
    """
    Generates a grid of evenly spaced points for visualizations
    returning grid of the first paramater, second parameter,
    and their values stacked with optionally fixed values.
    fixed -- contains position and argument to keep fixed
    """
    if not fixed:
        p1_low, p2_low = lower
        p1_high, p2_high = upper
    else:
        # find out which parameters are not fixed and get their bounds
        n_params = len(lower)

        p1_idx, p2_idx = [i for i in range(n_params)
                          if i not in [j for (j, _) in fixed]]

        p1_low, p2_low = lower[p1_idx], lower[p2_idx]
        p1_high, p2_high = upper[p1_idx], upper[p2_idx]

    p1_range = np.linspace(p1_low, p1_high, splits)
    p2_range = np.linspace(p2_low, p2_high, splits)
    p1_grid, p2_grid = np.meshgrid(p1_range, p2_range)

    if fixed:
        # create a grid for every parameter and insert in
        # corresponding position in the grids array
        grids = [None] * n_params
        grids[p1_idx] = p1_grid
        grids[p2_idx] = p2_grid
        for (i, val) in fixed:
            fixed_grid = np.zeros((splits, splits)) + val
            grids[i] = fixed_grid
    else:
        grids = [p1_grid, p2_grid]

    grid = np.dstack(tuple(grids))
    return p1_grid, p2_grid, grid


def predict_grid(model, grid):
    """
    Given a function and a grid of inputs this method
    calculates the output for each input in the grid.
    """
    rows, cols, n_params = grid.shape
    flatten_grid = grid.reshape((rows * cols, n_params))
    pred = np.apply_along_axis(model, 1, flatten_grid)
    return pred.reshape(rows, cols)
