####################################################################################################
###                                                                                ###
###                             Functions for simulating spectra                                 ###
###                               Author: Manuel Cordova (EPFL)                                  ###
###                                Last modified: 03.09.2021                                     ###
###                                                                                ###
####################################################################################################

# Import libraries
import numpy as np
import networkx as nx

# Set criterion for equivalent nodes
nm = nx.algorithms.isomorphism.categorical_node_match("elem", "X")

# Set criterion for equivalent edges
em = nx.algorithms.isomorphism.categorical_edge_match("w", -1)

def cleanup_methyl_protons(labels, Gs, envs, shifts, errs, ws, crysts, inds, atoms, bonds):
    """
    Gather methyl proton shifts into only one per methyl group

    Inputs:     - labels        List of labels for each distribution
                - Gs            Graph of each distribution
                - envs          Environment of each graph
                - shifts        Predicted shifts in each distribution
                - errs          Prediction errors in each distribution
                - ws            Maximum depth for each distribution
                - crysts        Crystals in each distribution
                - inds          Indices of the atoms in each distribution
                - atoms         List of atoms in the molecule
                - bonds         Bonded atoms for each atom in the molecule (by index)

    Outputs:    - new_labels    Cleaned list of labels for each distribution
                - new_Gs        Cleaned list of graphs
                - new_shifts    Cleaned list of predicted shifts in each distribution
                - new_errs      Cleaned list of prediction errors in each distribution
                - new_ws        Cleaned list of maximum depth for each distribution
                - new_crysts    Cleaned array of crystals in each distribution
                - new_inds      Cleaned array of indices of the atoms in each distribution
    """

    # Initialize new arrays
    new_labels = []
    new_Gs = []
    new_envs = []
    new_shifts = []
    new_errs = []
    new_ws = []
    new_crysts = []
    new_inds = []

    # Array to store already identified methyl groups
    methyls = []

    # Loop over all graphs
    for l, G, env, sh, er, w, cryst, ind in zip(labels, Gs, envs, shifts, errs, ws, crysts, inds):

        # Get the index of the central node
        i = G.nodes[0]["ind"]
        # Check that there is only one neighbour (H should be linked to only one atom
        if len(bonds[i]) == 1:
            j = bonds[i][0]

            # Check if there are at least three protons linked to the neighbour
            nH = [atoms[k] for k in bonds[j]].count("H")

            if nH >= 3:
                # If we find a new methyl proton, add it,
                #   otherwise skip it if it is part of an already known methyl
                if j not in methyls:
                    methyls.append(j)
                    new_labels.append(l)
                    new_Gs.append(G)
                    new_envs.append(env)
                    new_shifts.append(sh)
                    new_errs.append(er)
                    new_ws.append(w)
                    new_crysts.append(cryst)
                    new_inds.append(ind)

            # If this is not a methyl proton, add it
            else:
                new_labels.append(l)
                new_Gs.append(G)
                new_envs.append(env)
                new_shifts.append(sh)
                new_errs.append(er)
                new_ws.append(w)
                new_crysts.append(cryst)
                new_inds.append(ind)

        # If the proton is bonded to more than one atom, add it
        else:
            new_labels.append(l)
            new_Gs.append(G)
            new_envs.append(env)
            new_shifts.append(sh)
            new_errs.append(er)
            new_ws.append(w)
            new_crysts.append(cryst)
            new_inds.append(ind)

    return new_labels, new_Gs, new_envs, new_shifts, new_errs, new_ws, new_crysts, new_inds

def cleanup_methyls(labels, shifts, errs, ws, crysts, inds, hashes, atoms, bonds):
    """
    Gather methyl 2D shifts

    Inputs:     - labels        List of labels for each distribution
                - shifts        Predicted shifts in each distribution
                - errs          Prediction errors in each distribution
                - ws            Maximum depth for each distribution
                - crysts        Crystals in each distribution
                - inds          Indices of the atoms in each distribution
                - hashes        List of hashes for each graph
                - atoms         List of atoms in the molecule
                - bonds         Bonded atoms for each atom in the molecule (by index)

    Outputs:    - new_labels    Cleaned list of labels for each distribution
                - new_Gs        Cleaned list of graphs
                - new_shifts    Cleaned list of predicted shifts in each distribution
                - new_errs      Cleaned list of prediction errors in each distribution
                - new_ws        Cleaned list of maximum depth for each distribution
                - new_crysts    Cleaned array of crystals in each distribution
                - new_inds      Cleaned array of indices of the atoms in each distribution
    """

    # Initialize the updated lists
    new_labels = []
    new_shifts = []
    new_errs = []
    new_ws = []
    new_crysts = []
    new_inds = []
    new_hashes = []

    # Array to store already identified methyl groups
    methyls = []

    # Get element
    elem = ""
    for c in labels[0]:
        if c.isdigit():
            break
        elem += c

    # Loop over all distributions
    for l, sh, er, w, cryst, ind, h in zip(labels, shifts, errs, ws, crysts, inds, hashes):

        # Get the index of the central node
        ind = int(l.split("-")[0].split(elem)[1]) - 1
        i = [k for k, e in enumerate(atoms) if e == elem][ind]

        # Check if there are at least three protons linked to the atom
        nH = [atoms[k] for k in bonds[i]].count("H")

        if nH >= 3:

            # If we find a new methyl proton, add it,
            #   otherwise skip it if it is part of an already known methyl
            if i not in methyls:
                methyls.append(i)
                new_labels.append(l)
                new_shifts.append(sh)
                new_errs.append(er)
                new_ws.append(w)
                new_crysts.append(cryst)
                new_inds.append(ind)
                new_hashes.append(h)

        # If this is not a methyl, add it
        else:
            new_labels.append(l)
            new_shifts.append(sh)
            new_errs.append(er)
            new_ws.append(w)
            new_crysts.append(cryst)
            new_inds.append(ind)
            new_hashes.append(h)

    return new_labels, new_shifts, new_errs, new_ws, new_crysts, new_inds, new_hashes

def cleanup_equivalent(labels, shifts, errs, ws, crysts, inds, hashes):
    """
    Gather equivalent graphs (identified by their shift distributions)

    Inputs:     - labels        List of labels of the distributions
                - shifts        List of predicted shifts in each distribution
                - errs          List of predicted errors in each distribution
                - ws            List of weights of the distributions
                - crysts        List of crystals in each distribution
                - inds          List of the atoms in each distribution
                - hashes        List of hashes for each graph

    Outputs:    - new_labels    Updated list of labels of the distributions
                - new_shifts    Updated list of predicted shifts in each distribution
                - new_errs      Updated list of predicted errors in each distribution
                - new_ws        Updated list of weights of the distributions
                - new_crysts    Updated list of crystals in each distribution
                - new_inds      Updated list of the atoms in each distribution
    """

    # Initialize the updated lists
    new_labels = []
    new_shifts = []
    new_errs = []
    new_ws = []
    new_crysts = []
    new_inds = []
    new_hashes = []

    # Loop over all the distributions
    for l, sh, er, w, cr, ind, h in zip(labels, shifts, errs, ws, crysts, inds, hashes):

        # If the distribution is already found, modify the label
        if h in new_hashes:
            i = new_hashes.index(h)
            new_labels[i] += "/{}".format(l)

        # Otherwise, append the distribution to the updated list
        else:
            new_labels.append(l)
            new_shifts.append(sh)
            new_errs.append(er)
            new_ws.append(w)
            new_crysts.append(cr)
            new_inds.append(ind)
            new_hashes.append(h)

    return new_labels, new_shifts, new_errs, new_ws, new_crysts, new_inds, new_hashes

def get_lims_1D(all_shifts, all_errs, extend=0.1, dx="rms"):
    """
    Get the limits for a predicted 1D spectrum: obtain furthest peaks ± err to determine range, extend by a factor

    Inputs: - all_shifts    List of shifts in the distributions
            - all_errs      List of predicted errors in the distributions
            - extend        How much to extend the range (fraction of the original range, applied to both sides)
            - dx            How the error is incorporated to the minimum/maximum shift:
                                "sel": use the error of the minimum and maximum peaks
                                "max": use the maximum error
                                "mean": use the mean error
                                "rms": use the rms error

    Output: - lx            Limits in the x-dimension
    """

    # Initialize limits
    lx = np.ones(2) * np.mean(all_shifts[0])

    # Get minimum and maximum peak of each distribution
    for shifts, errs in zip(all_shifts, all_errs):

        # Get indices of the minimum and maximum peaks
        imin = np.argmin(shifts)
        imax = np.argmax(shifts)

        # Add the corresponding errors to the minimum and maximum peaks
        if dx == "sel":
            min_x = shifts[imin] - errs[imin]
            max_x = shifts[imax] + errs[imax]
        # Add the maximum error to the minimum and maximum peaks
        elif dx == "max":
            m = np.max(errs)
            min_x = shifts[imin] - m
            max_x = shifts[imax] + m
        # Add the mean error to the minimum and maximum peaks
        elif dx == "mean":
            m = np.mean(errs)
            min_x = shifts[imin] - m
            max_x = shifts[imax] + m
        # Add the rms error to the minimum and maximum peaks
        elif dx == "rms":
            m = np.sqrt(np.mean(np.square(errs)))
            min_x = shifts[imin] - m
            max_x = shifts[imax] + m
        else:
            raise ValueError("Unknown dx: {}".format(dx))

        # Get the limits
        lx[0] = min(lx[0], min_x)
        lx[1] = max(lx[1], max_x)

    # Get the range
    r = lx[1] - lx[0]

    # Extend the limits by a fraction of the range
    lx[0] -= extend * r
    lx[1] += extend * r

    return lx

def get_lims_2D(all_shifts, all_errs, extend=0.1, dx="rms"):
    """
    Get the limits for a predicted 1D spectrum: obtain furthest peaks ± err to determine range, extend by a factor

    Inputs:     - all_shifts    List of all shifts in the distributions
                - all_errs      List of predicted errors in the distributions
                - extend        How much to extend the range (fraction of the original range, applied to both sides)
                - dx            How the error is incorporated to the minimum/maximum shift:
                                    "sel": use the error of the minimum and maximum peaks
                                    "max": use the maximum error
                                    "mean": use the mean error
                                    "rms": use the rms error

    Outputs:    - lx            Limits in the x-dimension
                - ly            Limits in the y-dimension
    """

    # Initialize limits
    lx = np.ones(2) * np.mean(all_shifts[0][:,0])
    ly = np.ones(2) * np.mean(all_shifts[0][:,1])

    # Get minimum and maximum peak of each distribution in each dimension
    for shifts, errs in zip(all_shifts, all_errs):
        imin_x = np.argmin(shifts[:,0])
        imax_x = np.argmax(shifts[:,0])
        imin_y = np.argmin(shifts[:,1])
        imax_y = np.argmax(shifts[:,1])

        # Add the corresponding errors to the minimum and maximum peaks
        if dx == "sel":
            min_x = shifts[imin_x, 0] - errs[imin_x, 0]
            max_x = shifts[imax_x, 0] + errs[imax_x, 0]
            min_y = shifts[imin_y, 1] - errs[imin_y, 1]
            max_y = shifts[imax_y, 1] + errs[imax_y, 1]
        # Add the maximum error to the minimum and maximum peaks
        elif dx == "max":
            mx = np.max(errs[:, 0])
            my = np.max(errs[:, 1])
            min_x = shifts[imin_x, 0] - mx
            max_x = shifts[imax_x, 0] + mx
            min_y = shifts[imin_y, 1] - my
            max_y = shifts[imax_y, 1] + my
        # Add the mean error to the minimum and maximum peaks
        elif dx == "mean":
            mx = np.mean(errs[:, 0])
            my = np.mean(errs[:, 1])
            min_x = shifts[imin_x, 0] - mx
            max_x = shifts[imax_x, 0] + mx
            min_y = shifts[imin_y, 1] - my
            max_y = shifts[imax_y, 1] + my
        # Add the rms error to the minimum and maximum peaks
        elif dx == "rms":
            mx = np.sqrt(np.mean(np.square(errs[:, 0])))
            my = np.sqrt(np.mean(np.square(errs[:, 1])))
            min_x = shifts[imin_x, 0] - mx
            max_x = shifts[imax_x, 0] + mx
            min_y = shifts[imin_y, 1] - my
            max_y = shifts[imax_y, 1] + my
        else:
            raise ValueError("Unknown dx: {} (accepted values: 'sel', 'max', 'mean', 'rms')".format(dx))

        # Get the limits
        lx[0] = min(lx[0], min_x)
        lx[1] = max(lx[1], max_x)
        ly[0] = min(ly[0], min_y)
        ly[1] = max(ly[1], max_y)

        # Get the ranges
        rx = lx[1] - lx[0]
        ry = ly[1] - ly[0]

    # Extend the limits by a fraction of the ranges
    lx[0] -= extend * rx
    lx[1] += extend * rx
    ly[0] -= extend * ry
    ly[1] += extend * ry

    return np.array([lx, ly])

def make_1D_distribution(x, shifts, errs, norm=None, max_shifts=None, seed=None):
    """
    Generate 1D distribution of chemical shifts from an array of shifts and errors

    Inputs: - x             Points in the x-axis to draw the Gaussians on
            - shifts        List of shifts in the distribution
            - errs          List of predicted errors in the distribution
            - norm          Distribution normalization to apply
                                None: no normalization
                                "max": top of the distribution set to 1
            - max_shifts    Maximum number of shifts to consider when constructing the distribution
            - seed          Seed for the random selection of shifts

    Output: - y         Value of the distribution at each point of x
    """

    # Initialize y array
    y = np.zeros_like(x)

    # If there are too many shifts, randomly select a subset of length max_shifts
    if max_shifts is not None and max_shifts < len(shifts):
        if seed is not None:
            np.random.seed(seed)

        inds = np.random.choice(len(shifts), max_shifts, replace=False)

        # Add the Gaussians
        for x0, w in zip(shifts[inds], errs[inds]):
            y += 1. / (w * np.sqrt(2. * np.pi)) * np.exp(np.square(x - x0)/(-2. * np.square(w)))

    # Otherwise, use all shifts
    else:
        # Add the Gaussians
        for x0, w in zip(shifts, errs):
            y += 1. / (w * np.sqrt(2. * np.pi)) * np.exp(np.square(x - x0)/(-2. * np.square(w)))

    # Return the non-normalized sum
    if norm is None:
        return y
    # Normalize the maximum value to one
    elif norm == "max":
        return y/np.max(y)
    else:
        raise ValueError("Unknown normalization: {}".format(norm))

def make_1D_distributions(lims, n_points, all_shifts, all_errs, norm=None, max_shifts=None, seed=None):
    """
    Generate 1D distributions of chemical shifts from arrays of shifts and errors of each distribution

    Inputs:     - lims          Limits of the distributions
                - n_points      Number of points in the distributions
                - all_shifts    Array of shifts for each distribution
                - all_errs      Array of predicted error for each distribution
                - norm          Distribution normalization to apply
                                    None: no normalization
                                    "max": top of each distribution set to 1
                - max_shifts    Maximum number of shifts to consider when constructing the distribution
                - seed          Seed for the random selection of shifts

    Outputs:    - x             Array of shielding values to plot the distributions against
                - ys            List of distributions
    """

    # Construct the array of shielding values
    x = np.linspace(lims[0], lims[1], n_points)

    # Generate the distributions
    ys = []
    for i, (sh, er) in enumerate(zip(all_shifts, all_errs)):
        print("  Constructing distribution {}/{}...".format(i+1, len(all_shifts)))
        ys.append(make_1D_distribution(x, sh, er, norm=norm, max_shifts=max_shifts, seed=seed))
        print("  Distribution constructed!\n")

    return x, ys

def make_2D_distribution(x, y, shifts, errs, norm=None, max_shifts=None, seed=None):
    """

    Inputs: - x             Array of x-values to draw the Gaussians on
            - y             Array of y-values to draw the Gaussians on
            - shifts        List of shifts in the distribution
            - errs          List of predicted errors in the distribution
            - norm          Distribution normalization to apply
                                None: no normalization
                                "max": top of the distribution set to 1
            - max_shifts    Maximum number of shifts to consider when constructing the distribution
            - seed          Seed for the random selection of shifts

    Output: - Z             Value of the distribution at each point of the X-Y grid
    """

    # Initialize Z array
    Z = np.zeros((y.shape[0], x.shape[0]))

    # If there are too many shifts, randomly select a subset of length max_shifts
    if max_shifts is  not None and max_shifts < len(shifts):
        if seed is not None:
            np.random.seed(seed)

        inds = np.random.choice(len(shifts), max_shifts, replace=False)

        # Add the 2D Gaussians
        for [x0, y0], [wx, wy] in zip(shifts[inds], errs[inds]):
            gx = np.exp(np.square(x - x0) / (-2. * np.square(wx))) / wx
            gy = np.exp(np.square(y - y0) / (-2. * np.square(wy))) / wy
            Z += np.outer(gy, gx) / (2. * np.pi)

    # Otherwise, use all shifts
    else:
        # Add the 2D Gaussians
        for [x0, y0], [wx, wy] in zip(shifts, errs):
            gx = np.exp(np.square(x - x0) / (-2. * np.square(wx))) / wx
            gy = np.exp(np.square(y - y0) / (-2. * np.square(wy))) / wy
            Z += np.outer(gy, gx) / (2. * np.pi)

    if norm is None:
        return Z
    elif norm == "max":
        return Z / np.max(Z)
    else:
        raise ValueError("Unknown normalization: {}".format(norm))

def make_2D_distributions(lims, n_points, all_shifts, all_errs, norm=None, max_shifts=None, seed=None):
    """
    Generate 2D distributions of chemical shifts from arrays of shifts and errors of each distribution

    Inputs:     - lims          Limits of the distributions
                - n_points      Number of points in the distributions
                - all_shifts    Array of shifts for each distribution
                - all_errs      Array of predicted error for each distribution
                - norm          Distribution normalization to apply
                                    None: no normalization
                                    "max": top of each distribution set to 1
                - max_shifts    Maximum number of shifts to consider when constructing the distribution
                - seed          Seed for the random selection of shifts

    Outputs:    - X             Grid of shielding values (first dimension) to plot the distributions against
                - Y             Grid of shielding values (second dimension) to plot the distributions against
                - Zs            List of distributions
    """

    # Generate grid of X and Y values
    x = np.linspace(lims[0,0], lims[0,1], n_points)
    y = np.linspace(lims[1,0], lims[1,1], n_points)
    X, Y = np.meshgrid(x, y)

    # Generate the distributions
    Zs = []
    for i, (sh, er) in enumerate(zip(all_shifts, all_errs)):
        print("  Constructing distribution {}/{}...".format(i+1, len(all_shifts)))
        Zs.append(make_2D_distribution(x, y, sh, er, norm=norm, max_shifts=max_shifts, seed=seed))
        print("  Distribution constructed!\n")

    return X, Y, Zs

def get_distribution_max_1D(x, ys):
    """
    Obtain the maximum of each distribution

    Inputs: - x         Array of shielding values to plot the distributions against
            - ys        List of distributions

    Output: - centers   Maximum of each distribution
    """

    # Initialize array of centers
    centers = []

    # Get the center of each distribution (withing the set of shielding values considered)
    for i, y in enumerate(ys):
        inds = np.where(y == np.max(y))[0]

        if 0 in inds or (len(y) - 1) in inds:
            msg = "    WARNING: the maximum of distribution {} is at the edge of the"
            msg += " chemical shielding range! Consider expanding the range!".format(i+1)
            print(msg)

        centers.append(x[inds[0]])

    return np.array(centers)

def get_distribution_max_2D(X, Y, Zs):
    """
    Obtain the maximum of each distribution

    Inputs: - X         Grid of shielding values of the first element to plot the distribution against
            - Y         Grid of shielding values of the second element to plot the distribution against
            - Zs        List of distributions

    Output: - centers   Maximum of the distributions
    """

    #Initialize array of centers
    centers = []

    # Get the center of each distribution (within the set X and Y values)
    for i, Z in enumerate(Zs):

        inds_x, inds_y = np.where(Z == np.max(Z))

        if 0 in inds_x or 0 in inds_y or (Z.shape[0] - 1) in inds_x or (Z.shape[1] - 1) in inds_y:
            msg = "    WARNING: the maximum of distribution {} is at the edge of the"
            msg += " chemical shielding range! Consider expanding the range!".format(i+1)
            print(msg)

        centers.append([X[inds_x[0], inds_y[0]], Y[inds_x[0], inds_y[0]]])

    return np.array(centers)

def compute_scores_1D(exp, shifts, errs, conv, max_shifts=None, seed=None, acc=None, N=101):
    """
    Compute the scores for every possible assignment. If the variable "acc" is set to None, the probability density
    at the experimental shift yields the score. If "acc" is set to a value, the probability between
    the experimental shift e - acc and e + acc (computed as the numerical integral) is used as the score.

    Inputs: - exp           List of experimental shifts
            - shifts        List of shifts in each distribution
            - errs          List of errors in each distribution
            - conv          Conversion factors [slope, offset] from shielding to shift
            - max_shifts    Maximum number of shifts to select to construct the distribution
            - seed          Seed for random selection of shifts
            - acc           Accuracy of the shifts
            - N             Number of points in each shift if an accuracy is set

    Output: - scores    Matrix of scores for all possible assignments
    """

    # Initialize array of scores
    scores = np.zeros((len(shifts), len(exp)))

    # If no accuracy is set, take the shifts as elements of the array x
    if acc is None:
        x = np.array(exp)
    # Otherwise, append the array of N element between e - acc and e + acc to x, for each shift e
    else:
        x = []
        for e in exp:
            x.extend(list(np.linspace(e-acc, e+acc, N)))
        x = np.array(x)

    # Loop over all distributions
    for i, (sh, er) in enumerate(zip(shifts, errs)):
        print("  Evaluating distribution {}/{}...".format(i+1, len(shifts)))
        # Compute the values of the distribution on the array x
        y = make_1D_distribution(x, sh*conv[0]+conv[1], er, max_shifts=max_shifts, seed=seed)

        # If an accuracy is set, get the integral
        if acc is not None:
            y2 = []
            for j in range(len(exp)):
                y2.append(np.trapz(y[j*N:(j+1)*N], x=x[j*N:(j+1)*N]))
            y = np.array(y2)

        # Append the scores of this distribution
        if np.sum(y) < 1e-6:
            print("    WARNING: Distribution {} does not seem to match any experimental shift".format(i+1))
        scores[i] = y / np.sum(y)
        print("  Done!\n")

    return scores

def compute_scores_2D(exp, shifts, errs, conv_x, conv_y, max_shifts=None, seed=None, acc_x=None, acc_y=None, N=101):
    """
    Compute the scores for every possible assignment. If the variable "acc_x/acc_y" is set to None, the probability density
    at the experimental shift yields the score. If "acc_x/acc_y" are set to numerical values, the probability between
    the experimental shift e - acc and e + acc in each dimension (computed as the numerical integral) is used as the score.

    Inputs: - exp           List of experimental shifts
            - shifts        List of shifts in each distribution
            - errs          List of errors in each distribution
            - conv_x        Conversion factors [slope, offset] from shielding to shift in the x-axis
            - conv_y        Conversion factors [slope, offset] from shielding to shift in the y-axis
            - max_shifts    Maximum number of shifts to select to construct the distribution
            - seed          Seed for random selection of shifts
            - acc_x         Accuracy of the shifts in the x dimension
            - acc_y         Accuracy of the shifts in the y dimension
            - N             Number of points in each shift (along each axis) if an accuracy is set

    Output: - scores        Matrix of scores for all possible assignments
    """

    # Initialize matrix of scores
    scores = np.zeros((len(shifts), len(exp)))

    # If no accuracy is set, set x- and y-axes as the experimental shifts
    if acc_x is None and acc_y is None:
        x = np.array(exp)[:,0]
        y = np.array(exp)[:,1]

    # Accuracy set only in the x-axis
    elif acc_y is None:
        x = []
        for e in exp:
            x.extend(list(np.linspace(e[0]-acc_x, e[0]+acc_x, N)))
        x = np.array(x)
        y = np.array(exp)[:,1]

    # Accuracy set only in the y-axis
    elif acc_x is None:
        y = []
        for e in exp:
            y.extend(list(np.linspace(e[1]-acc_y, e[1]+acc_y, N)))
        y = np.array(y)
        x = np.array(exp)[:,0]

    # Accuracy set in both axes
    else:
        x = []
        y = []
        for e in exp:
            x.extend(list(np.linspace(e[0]-acc_x, e[0]+acc_x, N)))
            y.extend(list(np.linspace(e[1]-acc_y, e[1]+acc_y, N)))
        x = np.array(x)
        y = np.array(y)

    # Loop over all distributions
    for i, (s, er) in enumerate(zip(shifts, errs)):

        print("  Evaluating distribution {}/{}...".format(i+1, len(shifts)))

        conv_shifts = np.zeros_like(s)
        conv_shifts[:,0] = s[:,0]*conv_x[0]+conv_x[1]
        conv_shifts[:,1] = s[:,1]*conv_y[0]+conv_y[1]

        # If no accuracy is set, Compute the values of the distribution on the grid of experimental shifts
        if acc_x is None and acc_y is None:
            Z = make_2D_distribution(x, y, conv_shifts, er, max_shifts=max_shifts, seed=seed)
            these_scores = np.diag(Z)

        # If accuracy is set only along the x-axis, integrate over the range set
        elif acc_y is None:
            these_scores = np.zeros(len(exp))
            for j in range(len(exp)):
                Z = make_2D_distribution(x[j*N:(j+1)*N], y, conv_shifts, er, max_shifts=max_shifts, seed=seed)
                these_scores[j] = np.trapz(Z[j], x=x[j*N:(j+1)*N])

        # If accuracy is set only along the y-axis, integrate over the range set
        elif acc_x is None:
            these_scores = np.zeros(len(exp))
            for j in range(len(exp)):
                Z = make_2D_distribution(x, y[j*N:(j+1)*N], conv_shifts, er, max_shifts=max_shifts, seed=seed)
                these_scores[j] = np.trapz(Z[:, j], x=y[j*N:(j+1)*N])

        # If accuracy is set along both axes, integrate over the rectangle set
        else:
            these_scores = np.zeros(len(exp))
            for j in range(len(exp)):
                Z = make_2D_distribution(x[j*N:(j+1)*N], y[j*N:(j+1)*N], conv_shifts, er, max_shifts=max_shifts, seed=seed)
                these_scores[j] = np.trapz(np.trapz(Z, x=x[j*N:(j+1)*N]), x=y[j*N:(j+1)*N])

        if np.sum(these_scores) < 1e-6:
            print("    WARNING: Distribution {} does not seem to match any experimental shift".format(i+1))
        scores[i] = these_scores / np.sum(these_scores)
        print("  Done!\n")

    return scores


#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from __future__ import annotations

from collections import defaultdict
from functools import partial
from typing import (
    TYPE_CHECKING,
    Callable,
    Dict,
    List,
    MutableMapping,
    Optional,
    Tuple,
    Union,
)

import numpy as np
import torch
from ax.core.base_trial import TrialStatus
from ax.core.batch_trial import BatchTrial
from ax.core.experiment import Experiment
from ax.core.objective import MultiObjective, Objective, ScalarizedObjective
from ax.core.observation import Observation, ObservationData, ObservationFeatures
from ax.core.optimization_config import MultiObjectiveOptimizationConfig, TRefPoint
from ax.core.outcome_constraint import (
    ComparisonOp,
    OutcomeConstraint,
    ScalarizedOutcomeConstraint,
)
from ax.core.parameter import ChoiceParameter, ParameterType, RangeParameter
from ax.core.parameter_constraint import ParameterConstraint
from ax.core.search_space import SearchSpace, SearchSpaceDigest
from ax.core.trial import Trial
from ax.core.types import TBounds, TCandidateMetadata
from ax.modelbridge.transforms.base import Transform
from ax.models.torch.frontier_utils import (
    get_weighted_mc_objective_and_objective_thresholds,
    get_default_frontier_evaluator,
)
from ax.utils.common.logger import get_logger
from ax.utils.common.typeutils import (
    checked_cast,
    checked_cast_optional,
    not_none,
    checked_cast_to_tuple,
)
from botorch.utils.multi_objective.hypervolume import Hypervolume
from torch import Tensor

logger = get_logger(__name__)

if TYPE_CHECKING:
    # import as module to make sphinx-autodoc-typehints happy
    from ax import modelbridge as modelbridge_module  # noqa F401  # pragma: no cover

def extract_parameter_constraints(
    parameter_constraints: List[ParameterConstraint], param_names: List[str]
) -> Optional[TBounds]:
    """Extract parameter constraints."""
    if len(parameter_constraints) > 0:
        A = np.zeros((len(parameter_constraints), len(param_names)))
        b = np.zeros((len(parameter_constraints), 1))
        for i, c in enumerate(parameter_constraints):
            b[i, 0] = c.bound
            for name, val in c.constraint_dict.items():
                A[i, param_names.index(name)] = val
        linear_constraints: TBounds = (A, b)
    else:
        linear_constraints = None
    return linear_constraints

def extract_search_space_digest(
    search_space: SearchSpace, param_names: List[str]
) -> SearchSpaceDigest:
    """Extract basic parameter prpoerties from a search space."""
    bounds: List[Tuple[Union[int, float], Union[int, float]]] = []
    ordinal_features: List[int] = []
    categorical_features: List[int] = []
    discrete_choices: Dict[int, List[Union[int, float]]] = {}
    task_features: List[int] = []
    fidelity_features: List[int] = []
    target_fidelities: Dict[int, Union[int, float]] = {}

    for i, p_name in enumerate(param_names):
        p = search_space.parameters[p_name]
        if isinstance(p, ChoiceParameter):
            if p.is_task:
                task_features.append(i)
            elif p.is_ordered:
                ordinal_features.append(i)
            else:
                categorical_features.append(i)
            # at this point we can assume that values are numeric due to transforms
            discrete_choices[i] = p.values  # pyre-ignore [6]
            bounds.append((min(p.values), max(p.values)))  # pyre-ignore [6]
        elif isinstance(p, RangeParameter):
            if p.log_scale:
                raise ValueError(f"{p} is log scale")
            if p.parameter_type == ParameterType.INT:
                ordinal_features.append(i)
                d_choices = list(range(int(p.lower), int(p.upper) + 1))
                discrete_choices[i] = d_choices  # pyre-ignore [6]
            bounds.append((p.lower, p.upper))
        else:
            raise ValueError(f"Unknown parameter type {type(p)}")
        if p.is_fidelity:
            if not isinstance(not_none(p.target_value), (int, float)):
                raise NotImplementedError("Only numerical target values are supported.")
            target_fidelities[i] = checked_cast_to_tuple((int, float), p.target_value)
            fidelity_features.append(i)

    return SearchSpaceDigest(
        feature_names=param_names,
        bounds=bounds,
        ordinal_features=ordinal_features,
        categorical_features=categorical_features,
        discrete_choices=discrete_choices,
        task_features=task_features,
        fidelity_features=fidelity_features,
        target_fidelities=target_fidelities,
    )

def extract_objective_thresholds(
    objective_thresholds: TRefPoint,
    objective: Objective,
    outcomes: List[str],
) -> Optional[np.ndarray]:
    """Extracts objective thresholds' values, in the order of `outcomes`.

    Will return None if no objective thresholds, otherwise the extracted tensor
    will be the same length as `outcomes`.

    If one objective threshold is specified, they must be specified for every
    metric in the objective.

    Outcomes that are not part of an objective will be given a threshold of 0
    in this tensor, under the assumption that its value will not be used. Note
    that setting it to 0 for an outcome that is part of the objective would be
    incorrect, hence we validate that all objective metrics are represented.

    Args:
        objective_thresholds: Objective thresholds to extract values from.
        objective: The corresponding Objective, for validation purposes.
        outcomes: n-length list of names of metrics.

    Returns:
        (n,) array of thresholds
    """
    if len(objective_thresholds) == 0:
        return None

    objective_threshold_dict = {}
    for ot in objective_thresholds:
        if ot.relative:
            raise ValueError(
                f"Objective {ot.metric.name} has a relative threshold that is not "
                f"supported here."
            )
        objective_threshold_dict[ot.metric.name] = ot.bound

    if len(objective_threshold_dict) != len(objective.metrics):
        raise ValueError(
            "Objective thresholds do not match number of objective metrics."
        )
    # Initialize these to be nan to make sure that objective thresholds for
    # non-objective metrics are never used
    obj_t = np.full(len(outcomes), float("nan"))
    for metric in objective.metrics:
        if metric.name not in objective_threshold_dict:
            raise ValueError(
                f"Objective threshold not specified for {metric.name}. Thresholds must "
                f"be specified for all objective metrics or for none."
            )
        obj_t[outcomes.index(metric.name)] = objective_threshold_dict[metric.name]
    return obj_t

def extract_objective_weights(objective: Objective, outcomes: List[str]) -> np.ndarray:
    """Extract a weights for objectives.

    Weights are for a maximization problem.

    Give an objective weight to each modeled outcome. Outcomes that are modeled
    but not part of the objective get weight 0.

    In the single metric case, the objective is given either +/- 1, depending
    on the minimize flag.

    In the multiple metric case, each objective is given the input weight,
    multiplied by the minimize flag.

    Args:
        objective: Objective to extract weights from.
        outcomes: n-length list of names of metrics.

    Returns:
        n-length list of weights.

    """
    objective_weights = np.zeros(len(outcomes))
    if isinstance(objective, ScalarizedObjective):
        s = -1.0 if objective.minimize else 1.0
        for obj_metric, obj_weight in objective.metric_weights:
            objective_weights[outcomes.index(obj_metric.name)] = obj_weight * s
    elif isinstance(objective, MultiObjective):
        for obj, obj_weight in objective.objective_weights:
            s = -1.0 if obj.minimize else 1.0
            objective_weights[outcomes.index(obj.metric.name)] = obj_weight * s
    else:
        s = -1.0 if objective.minimize else 1.0
        objective_weights[outcomes.index(objective.metric.name)] = s
    return objective_weights

def extract_outcome_constraints(
    outcome_constraints: List[OutcomeConstraint], outcomes: List[str]
) -> TBounds:
    # Extract outcome constraints
    if len(outcome_constraints) > 0:
        A = np.zeros((len(outcome_constraints), len(outcomes)))
        b = np.zeros((len(outcome_constraints), 1))
        for i, c in enumerate(outcome_constraints):
            s = 1 if c.op == ComparisonOp.LEQ else -1
            if isinstance(c, ScalarizedOutcomeConstraint):
                for c_metric, c_weight in c.metric_weights:
                    j = outcomes.index(c_metric.name)
                    A[i, j] = s * c_weight
            else:
                j = outcomes.index(c.metric.name)
                A[i, j] = s
            b[i, 0] = s * c.bound
        outcome_constraint_bounds: TBounds = (A, b)
    else:
        outcome_constraint_bounds = None
    return outcome_constraint_bounds

def validate_and_apply_final_transform(
    objective_weights: np.ndarray,
    outcome_constraints: Optional[Tuple[np.ndarray, np.ndarray]],
    linear_constraints: Optional[Tuple[np.ndarray, np.ndarray]],
    pending_observations: Optional[List[np.ndarray]],
    objective_thresholds: Optional[np.ndarray] = None,
    final_transform: Callable[[np.ndarray], Tensor] = torch.tensor,
) -> Tuple[
    Tensor,
    Optional[Tuple[Tensor, Tensor]],
    Optional[Tuple[Tensor, Tensor]],
    Optional[List[Tensor]],
    Optional[Tensor],
]:
    # TODO: use some container down the road (similar to
    # SearchSpaceDigest) to limit the return arguments
    # pyre-fixme[35]: Target cannot be annotated.
    objective_weights: Tensor = final_transform(objective_weights)
    if outcome_constraints is not None:  # pragma: no cover
        # pyre-fixme[35]: Target cannot be annotated.
        outcome_constraints: Tuple[Tensor, Tensor] = (
            final_transform(outcome_constraints[0]),
            final_transform(outcome_constraints[1]),
        )
    if linear_constraints is not None:  # pragma: no cover
        # pyre-fixme[35]: Target cannot be annotated.
        linear_constraints: Tuple[Tensor, Tensor] = (
            final_transform(linear_constraints[0]),
            final_transform(linear_constraints[1]),
        )
    if pending_observations is not None:  # pragma: no cover
        # pyre-fixme[35]: Target cannot be annotated.
        pending_observations: List[Tensor] = [
            final_transform(pending_obs) for pending_obs in pending_observations
        ]
    if objective_thresholds is not None:
        # pyre-fixme[35]: Target cannot be annotated.
        objective_thresholds: Tensor = final_transform(objective_thresholds)
    return (
        objective_weights,
        outcome_constraints,
        linear_constraints,
        pending_observations,
        objective_thresholds,
    )

def get_fixed_features(
    fixed_features: ObservationFeatures, param_names: List[str]
) -> Optional[Dict[int, float]]:
    """Reformat a set of fixed_features."""
    fixed_features_dict = {}
    for p_name, val in fixed_features.parameters.items():
        # These all need to be floats at this point.
        # pyre-ignore[6]: All float here.
        val_ = float(val)
        fixed_features_dict[param_names.index(p_name)] = val_
    fixed_features_dict = fixed_features_dict if len(fixed_features_dict) > 0 else None
    return fixed_features_dict

def pending_observations_as_array(
    pending_observations: Dict[str, List[ObservationFeatures]],
    outcome_names: List[str],
    param_names: List[str],
) -> Optional[List[np.ndarray]]:
    """Re-format pending observations.

    Args:
        pending_observations: List of raw numpy pending observations.
        outcome_names: List of outcome names.
        param_names: List fitted param names.

    Returns:
        Filtered pending observations data, by outcome and param names.
    """
    if len(pending_observations) == 0:
        pending_array: Optional[List[np.ndarray]] = None
    else:
        pending_array = [np.array([]) for _ in outcome_names]
        for metric_name, po_list in pending_observations.items():
            # It is possible that some metrics attached to the experiment should
            # not be included in pending features for a given model. For example,
            # if a model is fit to the initial data that is missing some of the
            # metrics on the experiment or if a model just should not be fit for
            # some of the metrics attached to the experiment, so metrics that
            # appear in pending_observations (drawn from an experiment) but not
            # in outcome_names (metrics, expected for the model) are filtered out.ß
            if metric_name not in outcome_names:
                continue
            pending_array[outcome_names.index(metric_name)] = np.array(
                [[po.parameters[p] for p in param_names] for po in po_list]
            )
    return pending_array

def parse_observation_features(
    X: np.ndarray,
    param_names: List[str],
    candidate_metadata: Optional[List[TCandidateMetadata]] = None,
) -> List[ObservationFeatures]:
    """Re-format raw model-generated candidates into ObservationFeatures.

    Args:
        param_names: List of param names.
        X: Raw np.ndarray of candidate values.
        candidate_metadata: Model's metadata for candidates it produced.

    Returns:
        List of candidates, represented as ObservationFeatures.
    """
    if candidate_metadata and len(candidate_metadata) != len(X):
        raise ValueError(  # pragma: no cover
            "Observations metadata list provided is not of "
            "the same size as the number of candidates."
        )
    observation_features = []
    for i, x in enumerate(X):
        observation_features.append(
            ObservationFeatures(
                parameters=dict(zip(param_names, x)),
                metadata=candidate_metadata[i] if candidate_metadata else None,
            )
        )
    return observation_features

def transform_callback(
    param_names: List[str], transforms: MutableMapping[str, Transform]
) -> Callable[[np.ndarray], np.ndarray]:
    """A closure for performing the `round trip` transformations.

    The function round points by de-transforming points back into
    the original space (done by applying transforms in reverse), and then
    re-transforming them.
    This function is specifically for points which are formatted as numpy
    arrays. This function is passed to _model_gen.

    Args:
        param_names: Names of parameters to transform.
        transforms: Ordered set of transforms which were applied to the points.

    Returns:
        a function with for performing the roundtrip transform.
    """

    def _roundtrip_transform(x: np.ndarray) -> np.ndarray:
        """Inner function for performing aforementioned functionality.

        Args:
            x: points in the transformed space (e.g. all transforms have been applied
                to them)

        Returns:
            points in the transformed space, but rounded via the original space.
        """
        # apply reverse terminal transform to turn array to ObservationFeatures
        observation_features = [
            ObservationFeatures(
                parameters={p: float(x[i]) for i, p in enumerate(param_names)}
            )
        ]
        # reverse loop through the transforms and do untransform
        for t in reversed(transforms.values()):
            observation_features = t.untransform_observation_features(
                observation_features
            )
        # forward loop through the transforms and do transform
        for t in transforms.values():
            observation_features = t.transform_observation_features(
                observation_features
            )
        # parameters are guaranteed to be float compatible here, but pyre doesn't know
        new_x: List[float] = [
            # pyre-fixme[6]: Expected `Union[_SupportsIndex, bytearray, bytes, str,
            #  typing.SupportsFloat]` for 1st param but got `Union[None, bool, float,
            #  int, str]`.
            float(observation_features[0].parameters[p])
            for p in param_names
        ]
        # turn it back into an array
        return np.array(new_x)

    return _roundtrip_transform

def get_pending_observation_features(
    experiment: Experiment, include_failed_as_pending: bool = False
) -> Optional[Dict[str, List[ObservationFeatures]]]:
    """Computes a list of pending observation features (corresponding to arms that
    have been generated and deployed in the course of the experiment, but have not
    been completed with data or to arms that have been abandoned or belong to
    abandoned trials).

    NOTE: Pending observation features are passed to the model to
    instruct it to not generate the same points again.

    Args:
        experiment: Experiment, pending features on which we seek to compute.
        include_failed_as_pending: Whether to include failed trials as pending
            (for example, to avoid the model suggesting them again).

    Returns:
        An optional mapping from metric names to a list of observation features,
        pending for that metric (i.e. do not have evaluation data for that metric).
        If there are no pending features for any of the metrics, return is None.
    """
    pending_features = {}
    # Note that this assumes that if a metric appears in fetched data, the trial is
    # not pending for the metric. Where only the most recent data matters, this will
    # work, but may need to add logic to check previously added data objects, too.
    for trial_index, trial in experiment.trials.items():
        dat = trial.lookup_data()
        for metric_name in experiment.metrics:
            if metric_name not in pending_features:
                pending_features[metric_name] = []
            include_since_failed = include_failed_as_pending and trial.status.is_failed
            if isinstance(trial, BatchTrial):
                if trial.status.is_abandoned or (
                    (trial.status.is_deployed or include_since_failed)
                    and metric_name not in dat.df.metric_name.values
                    and trial.arms is not None
                ):
                    for arm in trial.arms:
                        not_none(pending_features.get(metric_name)).append(
                            ObservationFeatures.from_arm(
                                arm=arm, trial_index=np.int64(trial_index)
                            )
                        )
                abandoned_arms = trial.abandoned_arms
                for abandoned_arm in abandoned_arms:
                    not_none(pending_features.get(metric_name)).append(
                        ObservationFeatures.from_arm(
                            arm=abandoned_arm, trial_index=np.int64(trial_index)
                        )
                    )

            if isinstance(trial, Trial):
                if trial.status.is_abandoned or (
                    (trial.status.is_deployed or include_since_failed)
                    and metric_name not in dat.df.metric_name.values
                    and trial.arm is not None
                ):
                    not_none(pending_features.get(metric_name)).append(
                        ObservationFeatures.from_arm(
                            arm=not_none(trial.arm), trial_index=np.int64(trial_index)
                        )
                    )
    return pending_features if any(x for x in pending_features.values()) else None

def get_pending_observation_features_based_on_trial_status(
    experiment: Experiment,
) -> Optional[Dict[str, List[ObservationFeatures]]]:
    """A faster analogue of ``get_pending_observation_features`` that makes
    assumptions about trials in experiment in order to speed up extraction
    of pending points.

    Assumptions:

    * All arms in all trials in ``STAGED,`` ``RUNNING`` and ``ABANDONED`` statuses
      are to be considered pending for all outcomes.
    * All arms in all trials in other statuses are to be considered not pending for
      all outcomes.

    This entails:

    * No actual data-fetching for trials to determine whether arms in them are pending
      for specific outcomes.
    * Even if data is present for some outcomes in ``RUNNING`` trials, their arms will
      still be considered pending for those outcomes.

    NOTE: This function should not be used to extract pending features in field
    experiments, where arms in running trials should not be considered pending if
    there is data for those arms.

    Args:
        experiment: Experiment, pending features on which we seek to compute.

    Returns:
        An optional mapping from metric names to a list of observation features,
        pending for that metric (i.e. do not have evaluation data for that metric).
        If there are no pending features for any of the metrics, return is None.
    """
    pending_features = defaultdict(list)
    for status in [TrialStatus.STAGED, TrialStatus.RUNNING, TrialStatus.ABANDONED]:
        for trial in experiment.trials_by_status[status]:
            for metric_name in experiment.metrics:
                pending_features[metric_name].extend(
                    ObservationFeatures.from_arm(
                        arm=arm, trial_index=np.int64(trial.index)
                    )
                    for arm in trial.arms
                )

    return dict(pending_features) if any(x for x in pending_features.values()) else None

def clamp_observation_features(
    observation_features: List[ObservationFeatures], search_space: SearchSpace
) -> List[ObservationFeatures]:
    range_parameters = [
        p for p in search_space.parameters.values() if isinstance(p, RangeParameter)
    ]
    for obsf in observation_features:
        for p in range_parameters:
            if p.name not in obsf.parameters:
                continue
            if p.parameter_type == ParameterType.FLOAT:
                val = checked_cast(float, obsf.parameters[p.name])
            else:
                val = checked_cast(int, obsf.parameters[p.name])
            if val < p.lower:
                logger.info(
                    f"Untransformed parameter {val} "
                    f"less than lower bound {p.lower}, clamping"
                )
                obsf.parameters[p.name] = p.lower
            elif val > p.upper:
                logger.info(
                    f"Untransformed parameter {val} "
                    f"greater than upper bound {p.upper}, clamping"
                )
                obsf.parameters[p.name] = p.upper
    return observation_features

def get_pareto_frontier_and_transformed_configs(
    modelbridge: modelbridge_module.array.ArrayModelBridge,
    observation_features: List[ObservationFeatures],
    observation_data: Optional[List[ObservationData]] = None,
    objective_thresholds: Optional[TRefPoint] = None,
    optimization_config: Optional[MultiObjectiveOptimizationConfig] = None,
    arm_names: Optional[List[Optional[str]]] = None,
    use_model_predictions: bool = True,
) -> Tuple[List[Observation], Tensor, Tensor, Optional[Tensor]]:
    """Helper that applies transforms and calls frontier_evaluator.

    Returns transformed configs in addition to the Pareto observations.

    Args:
        modelbridge: Modelbridge used to predict metrics outcomes.
        observation_features: observation features to predict, if provided and
            use_model_predictions is True.
        observation_data: data for computing the Pareto front, unless features
            are provided and model_predictions is True.
        objective_thresholds: metric values bounding the region of interest in
            the objective outcome space.
        optimization_config: Optimization config.
        arm_names: Arm names for each observation.
        use_model_predictions: If True, will use model predictions at
            observation_features to compute Pareto front, if provided. If False,
            will use observation_data directly to compute Pareto front, regardless
            of whether observation_features are provided.

    Returns:
        frontier_observations: Observations of points on the pareto frontier.
        f: n x m tensor representation of the Pareto frontier values where n is the
        length of frontier_observations and m is the number of metrics.
        obj_w: m tensor of objective weights.
        obj_t: m tensor of objective thresholds corresponding to Y, or None if no
        objective thresholds used.
    """

    array_to_tensor = partial(_array_to_tensor, modelbridge=modelbridge)
    X = (
        modelbridge.transform_observation_features(observation_features)
        if use_model_predictions
        else None
    )
    X = array_to_tensor(X) if X is not None else None
    Y, Yvar = (None, None)
    if observation_data is not None:
        Y, Yvar = modelbridge.transform_observation_data(observation_data)
        Y, Yvar = (array_to_tensor(Y), array_to_tensor(Yvar))
    if arm_names is None:
        arm_names = [None] * len(observation_features)

    # Optimization_config
    mooc = optimization_config or checked_cast_optional(
        MultiObjectiveOptimizationConfig, modelbridge._optimization_config
    )
    if not mooc:
        raise ValueError(
            (
                "Experiment must have an existing optimization_config "
                "of type `MultiObjectiveOptimizationConfig` "
                "or `optimization_config` must be passed as an argument."
            )
        )
    if not isinstance(mooc, MultiObjectiveOptimizationConfig):
        mooc = not_none(MultiObjectiveOptimizationConfig.from_opt_conf(mooc))
    if objective_thresholds:
        mooc = mooc.clone_with_args(objective_thresholds=objective_thresholds)

    optimization_config = mooc

    # Transform OptimizationConfig.
    optimization_config = modelbridge.transform_optimization_config(
        optimization_config=optimization_config,
        fixed_features=ObservationFeatures(parameters={}),
    )
    # Extract weights, constraints, and objective_thresholds
    objective_weights = extract_objective_weights(
        objective=optimization_config.objective, outcomes=modelbridge.outcomes
    )
    outcome_constraints = extract_outcome_constraints(
        outcome_constraints=optimization_config.outcome_constraints,
        outcomes=modelbridge.outcomes,
    )
    obj_t = extract_objective_thresholds(
        objective_thresholds=optimization_config.objective_thresholds,
        objective=optimization_config.objective,
        outcomes=modelbridge.outcomes,
    )
    obj_t = array_to_tensor(obj_t)
    # Transform to tensors.
    obj_w, oc_c, _, _, _ = validate_and_apply_final_transform(
        objective_weights=objective_weights,
        outcome_constraints=outcome_constraints,
        linear_constraints=None,
        pending_observations=None,
        final_transform=array_to_tensor,
    )
    frontier_evaluator = get_default_frontier_evaluator()
    # pyre-ignore[28]: Unexpected keyword `modelbridge` to anonymous call
    f, cov, indx = frontier_evaluator(
        model=modelbridge.model,
        X=X,
        Y=Y,
        Yvar=Yvar,
        objective_thresholds=obj_t,
        objective_weights=obj_w,
        outcome_constraints=oc_c,
    )
    f, cov = f.detach().cpu().clone(), cov.detach().cpu().clone()
    indx = indx.tolist()
    frontier_observation_data = array_to_observation_data(
        f=f.numpy(), cov=cov.numpy(), outcomes=not_none(modelbridge.outcomes)
    )
    # Untransform observations
    for t in reversed(modelbridge.transforms.values()):  # noqa T484
        frontier_observation_data = t.untransform_observation_data(
            frontier_observation_data, []
        )
    # Construct observations
    frontier_observations = []
    for i, obsd in enumerate(frontier_observation_data):
        frontier_observations.append(
            Observation(
                features=observation_features[indx[i]],
                data=obsd,
                arm_name=arm_names[indx[i]],
            )
        )
    return frontier_observations, f, obj_w.cpu(), obj_t.cpu()

def pareto_frontier(
    modelbridge: modelbridge_module.array.ArrayModelBridge,
    observation_features: List[ObservationFeatures],
    observation_data: Optional[List[ObservationData]] = None,
    objective_thresholds: Optional[TRefPoint] = None,
    optimization_config: Optional[MultiObjectiveOptimizationConfig] = None,
    arm_names: Optional[List[Optional[str]]] = None,
    use_model_predictions: bool = True,
) -> List[Observation]:
    """Helper that applies transforms and calls frontier_evaluator.

    Args:
        modelbridge: Modelbridge used to predict metrics outcomes.
        observation_features: observation features to predict, if provided and
            use_model_predictions is True.
        observation_data: data for computing the Pareto front, unless features
            are provided and model_predictions is True.
        objective_thresholds: metric values bounding the region of interest in
            the objective outcome space.
        optimization_config: Optimization config.
        arm_names: Arm names for each observation.
        use_model_predictions: If True, will use model predictions at
            observation_features to compute Pareto front, if provided. If False,
            will use observation_data directly to compute Pareto front, regardless
            of whether observation_features are provided.

    Returns:
        frontier_observations: Observations of points on the pareto frontier.
    """
    return get_pareto_frontier_and_transformed_configs(
        modelbridge=modelbridge,
        observation_features=observation_features,
        observation_data=observation_data,
        objective_thresholds=objective_thresholds,
        optimization_config=optimization_config,
        arm_names=arm_names,
        use_model_predictions=use_model_predictions,
    )[0]

def predicted_pareto_frontier(
    modelbridge: modelbridge_module.array.ArrayModelBridge,
    objective_thresholds: Optional[TRefPoint] = None,
    observation_features: Optional[List[ObservationFeatures]] = None,
    optimization_config: Optional[MultiObjectiveOptimizationConfig] = None,
) -> List[Observation]:
    """Generate a pareto frontier based on the posterior means of given
    observation features.

    Given a model and features to evaluate use the model to predict which points
    lie on the pareto frontier.

    Args:
        modelbridge: Modelbridge used to predict metrics outcomes.
        objective_thresholds: metric values bounding the region of interest in
            the objective outcome space.
        observation_features: observation features to predict. Model's training
            data used by default if unspecified.
        optimization_config: Optimization config

    Returns:
        Observations representing points on the pareto frontier.
    """
    if observation_features is None:
        observation_features = []
        arm_names = []
        for obs in modelbridge.get_training_data():
            observation_features.append(obs.features)
            arm_names.append(obs.arm_name)
    else:
        arm_names = None
    if not observation_features:
        raise ValueError(
            "Must receive observation_features as input or the model must "
            "have training data."
        )

    pareto_observations = pareto_frontier(
        modelbridge=modelbridge,
        objective_thresholds=objective_thresholds,
        observation_features=observation_features,
        optimization_config=optimization_config,
        arm_names=arm_names,
    )
    return pareto_observations

def observed_pareto_frontier(
    modelbridge: modelbridge_module.array.ArrayModelBridge,
    objective_thresholds: Optional[TRefPoint] = None,
    optimization_config: Optional[MultiObjectiveOptimizationConfig] = None,
) -> List[Observation]:
    """Generate a pareto frontier based on observed data.

    Given observed data, return those outcomes in the pareto frontier.

    Args:
        modelbridge: Modelbridge that holds previous training data.
        objective_thresholds: metric values bounding the region of interest in
            the objective outcome space.
        optimization_config: Optimization config

    Returns:
        Data representing points on the pareto frontier.
    """
    # Get observation_data from current training data
    observation_data = []
    observation_features = []
    arm_names = []
    for obs in modelbridge.get_training_data():
        observation_data.append(obs.data)
        observation_features.append(obs.features)
        arm_names.append(obs.arm_name)

    pareto_observations = pareto_frontier(
        modelbridge=modelbridge,
        objective_thresholds=objective_thresholds,
        observation_data=observation_data,
        observation_features=observation_features,
        optimization_config=optimization_config,
        arm_names=arm_names,
        use_model_predictions=False,
    )
    return pareto_observations

def hypervolume(
    modelbridge: modelbridge_module.array.ArrayModelBridge,
    observation_features: List[ObservationFeatures],
    objective_thresholds: Optional[TRefPoint] = None,
    observation_data: Optional[List[ObservationData]] = None,
    optimization_config: Optional[MultiObjectiveOptimizationConfig] = None,
    use_model_predictions: bool = True,
) -> float:
    """Helper function that computes hypervolume of a given list of outcomes."""
    # Get Pareto front
    observations, f, obj_w, obj_t = get_pareto_frontier_and_transformed_configs(
        modelbridge=modelbridge,
        objective_thresholds=objective_thresholds,
        observation_features=observation_features,
        observation_data=observation_data,
        optimization_config=optimization_config,
        use_model_predictions=use_model_predictions,
    )
    if obj_t is None:
        raise ValueError(
            "Cannot compute hypervolume without having objective thresholds specified."
        )
    # Apply appropriate weights and thresholds
    obj, obj_t = get_weighted_mc_objective_and_objective_thresholds(
        objective_weights=obj_w, objective_thresholds=obj_t
    )
    f_t = obj(f)
    hv = Hypervolume(ref_point=obj_t)
    return hv.compute(f_t)

def predicted_hypervolume(
    modelbridge: modelbridge_module.array.ArrayModelBridge,
    objective_thresholds: Optional[TRefPoint] = None,
    observation_features: Optional[List[ObservationFeatures]] = None,
    optimization_config: Optional[MultiObjectiveOptimizationConfig] = None,
) -> float:
    """Calculate hypervolume of a pareto frontier based on the posterior means of
    given observation features.

    Given a model and features to evaluate calculate the hypervolume of the pareto
    frontier formed from their predicted outcomes.

    Args:
        modelbridge: Modelbridge used to predict metrics outcomes.
        objective_thresholds: point defining the origin of hyperrectangles that
            can contribute to hypervolume.
        observation_features: observation features to predict. Model's training
            data used by default if unspecified.
        optimization_config: Optimization config

    Returns:
        calculated hypervolume.
    """
    observation_features = (
        observation_features
        if observation_features is not None
        else [obs.features for obs in modelbridge.get_training_data()]
    )
    if not observation_features:
        raise ValueError(
            "Must receive observation_features as input or the model must "
            "have training data."
        )

    return hypervolume(
        modelbridge=modelbridge,
        objective_thresholds=objective_thresholds,
        observation_features=observation_features,
        optimization_config=optimization_config,
    )

def observed_hypervolume(
    modelbridge: modelbridge_module.array.ArrayModelBridge,
    objective_thresholds: Optional[TRefPoint] = None,
    optimization_config: Optional[MultiObjectiveOptimizationConfig] = None,
) -> float:
    """Calculate hypervolume of a pareto frontier based on observed data.

    Given observed data, return the hypervolume of the pareto frontier formed from
    those outcomes.

    Args:
        modelbridge: Modelbridge that holds previous training data.
        objective_thresholds: point defining the origin of hyperrectangles that
            can contribute to hypervolume.
        observation_features: observation features to predict. Model's training
            data used by default if unspecified.
        optimization_config: Optimization config

    Returns:
        (float) calculated hypervolume.
    """
    # Get observation_data from current training data.
    observation_data = [obs.data for obs in modelbridge.get_training_data()]
    observation_features = [obs.features for obs in modelbridge.get_training_data()]

    return hypervolume(
        modelbridge=modelbridge,
        objective_thresholds=objective_thresholds,
        observation_features=observation_features,
        observation_data=observation_data,
        optimization_config=optimization_config,
        use_model_predictions=False,
    )

def array_to_observation_data(
    f: np.ndarray, cov: np.ndarray, outcomes: List[str]
) -> List[ObservationData]:
    """Convert arrays of model predictions to a list of ObservationData.

    Args:
        f: An (n x m) array
        cov: An (n x m x m) array
        outcomes: A list of d outcome names

    Returns: A list of n ObservationData
    """
    observation_data = []
    for i in range(f.shape[0]):
        observation_data.append(
            ObservationData(
                metric_names=list(outcomes),
                means=f[i, :].copy(),
                covariance=cov[i, :, :].copy(),
            )
        )
    return observation_data

def observation_data_to_array(
    outcomes: List[str],
    observation_data: List[ObservationData],
) -> Tuple[np.ndarray, np.ndarray]:
    """Convert a list of Observation data to arrays.

    Args:
        observation_data: A list of n ObservationData

    Returns:
        An array of n ObservationData, each containing
            - f: An (n x m) array
            - cov: An (n x m x m) array
    """
    means = []
    cov = []
    for obsd in observation_data:
        metric_idxs = np.array([obsd.metric_names.index(m) for m in outcomes])
        means.append(obsd.means[metric_idxs])
        cov.append(obsd.covariance[metric_idxs][:, metric_idxs])
    return np.array(means), np.array(cov)

def observation_features_to_array(
    parameters: List[str], obsf: List[ObservationFeatures]
) -> np.ndarray:
    """Convert a list of Observation features to arrays."""
    return np.array([[of.parameters[p] for p in parameters] for of in obsf])

def _array_to_tensor(
    array: Union[np.ndarray, List[float]],
    modelbridge: Optional[modelbridge_module.base.ModelBridge] = None,
) -> Tensor:
    if modelbridge and hasattr(modelbridge, "_array_to_tensor"):
        # pyre-ignore[16]: modelbridge does not have attribute `_array_to_tensor`
        return modelbridge._array_to_tensor(array)
    else:
        return torch.tensor(array)


"""Module to hold our forward models that until now have been floating in the
web application.

Right now this includes database interaction and model rescaling software.

Authors
-------
Jules Fowler, April 2019
Natasha Batalha
Hannah Wakeford

Use
---
This is meant to be run through the Flask application, or can be run manually
with a provided list or arguments.
"""

## -- IMPORTS
import os

import astropy.constants as constants
from six.moves import StringIO
import astropy.table as at
import astropy.units as u
from bokeh.resources import INLINE
from bokeh.util.string import encode_utf8
from bokeh.embed import components
from bokeh.models import Range1d
from bokeh.models.widgets import Panel, Tabs
from bokeh.plotting import figure, output_file, save
import h5py
import numpy as np
import pandas as pd
from sqlalchemy import create_engine

from exoctk.utils import get_env_variables

## -- FUNCTIONS

def fortney_grid(args, write_plot=False, write_table=False):
    """
    Function to grab a Fortney Grid model, plot it, and make a table.

    Parameters
    ----------
    args : dict
        Dictionary of arguments for the Fortney Grid. Must include :
        temp
        chem
        cloud
        pmass
        m_unit
        reference_radius
        r_unit
        rstar
        rstar_unit
    write_plot : bool, optional
        Whether or not to save the bokeh plot, defaults to False.
    write_table : bool, optional
        Whether or not to save the ascii table, defaults to False.

    Returns
    -------
    fig : bokeh object
        The unsaved bokeh plot.
    fh : ascii table object
        The unsaved ascii table.
    temp_out : list of str of int
        The list of temperatures in the model grid.
    """

    # Check for Fortney Grid database
    print(os.path.join(get_env_variables()['exoctk_data'], 'fortney/fortney_models.db'))
    try:
        db = create_engine('sqlite:///' +
                os.path.join(get_env_variables()['exoctk_data'], 'fortney/fortney_models.db'))
        header = pd.read_sql_table('header', db)
    except:
        raise Exception('Fortney Grid File Path is incorrect, or not initialized')

    if args:
        rstar = float(args['rstar'])
        rstar = (rstar * u.Unit(args['rstar_unit'])).to(u.km)
        reference_radius = float(args['reference_radius'])
        rplan = (reference_radius * u.Unit(args['r_unit'])).to(u.km)
        temp = float(args['temp'])
        # clouds
        cloud = args['cloud']
        if cloud.find('flat') != -1:
            flat = int(cloud[4:])
            ray = 0
        elif cloud.find('ray') != -1:
            ray = int(cloud[3:])
            flat = 0
        elif int(cloud) == 0:
            flat = 0
            ray = 0
        else:
            flat = 0
            ray = 0
            print('No cloud parameter not specified, default no clouds added')

        # chemistry
        chem = args['chem']
        if chem == 'noTiO':
            noTiO = True
        if chem == 'eqchem':
            noTiO = False
            # grid does not allow clouds for cases with TiO
            flat = 0
            ray = 0

        fort_grav = 25.0 * u.m / u.s**2

        df = header.loc[(header.gravity == fort_grav) & (header.temp == temp) &
                        (header.noTiO == noTiO) & (header.ray == ray) &
                        (header.flat == flat)]

        wave_planet = np.array(pd.read_sql_table(df['name'].values[0], db)['wavelength'])[::-1]
        r_lambda = np.array(pd.read_sql_table(df['name'].values[0], db)['radius']) * u.km

        # All fortney models have fixed 1.25 radii
        z_lambda = r_lambda - (1.25 * u.R_jup).to(u.km)

        # Scale with planetary mass
        pmass = float(args['pmass'])
        mass = (pmass * u.Unit(args['m_unit'])).to(u.kg)

        # Convert radius to m for gravity units
        gravity = constants.G * (mass) / (rplan.to(u.m))**2.0

        # Scale lambbda (this technically ignores the fact that scaleheight
        # is altitude dependent) therefore, it will not be valide for very
        # very low gravities
        z_lambda = z_lambda * fort_grav / gravity

        # Create new wavelength dependent R based on scaled ravity
        r_lambda = z_lambda + rplan

        # Finally compute (rp/r*)^2
        flux_planet = np.array(r_lambda**2 / rstar**2)

        x = wave_planet
        y = flux_planet[::-1]

    else:
        df = pd.read_sql_table('t1000g25_noTiO', db)
        x, y = df['wavelength'], df['radius']**2.0 / 7e5**2.0

    tab = at.Table(data=[x, y])
    fh = StringIO()
    tab.write(fh, format='ascii.no_header')

    if write_table:
        tab.write('fortney.dat', format='ascii.no_header')

    fig = figure(plot_width=1100, plot_height=400)
    fig.line(x, 1e6 * (y - np.mean(y)), color='Black', line_width=0.5)
    fig.xaxis.axis_label = 'Wavelength (um)'
    fig.yaxis.axis_label = 'Rel. Transit Depth (ppm)'

    if write_plot:
        output_file('fortney.html')
        save(fig)

    # Return temperature list for the fortney grid page
    temp_out = list(map(str, header.temp.unique()))

    return fig, fh, temp_out

def generic_grid(input_args, write_plot=False, write_table=False):
    """
    Build a plot and table from the generic grid results.

    Parameters
    ----------
    input_args : dict
        A dictionary of the form output from the generic grid form.
        If manual input must include :
        r_star : The radius of the star.
        r_planet : The radius of the planet.
        gravity : The gravity.
        temperature : The temperature.
        condensation : local or rainout
        metallicity
        c_o : carbon/oxygen ratio
        haze
        cloud
    write_plot : bool, optional
        Whether to write the plot out. Defaults to False.
    write_table : bool, optional
        Whether to write the table out. Defaults to Fals.

    Returns
    -------
    plot : bokeh object
        Unsaved bokeh plot.
    table : ascii table object
        Unsaved ascii table.
    closest_match : dict
        A dictionary with the parameters/model name of the closest
        match in the grid.
    error_message : str
        An error message, or lack therof.
    """

    try:
        str_args = {}
        for key, value in input_args.items():
            str_args[key] = value[0]
        input_args=str_args
    except:
        pass
        #this attempts to convert passed arguments from list to string, and does nothing if the arguments are not lists

    # Find path to the database.
    database_path = os.path.join(get_env_variables()['exoctk_data'], 'generic/generic_grid_db.hdf5')
    # Build rescaled model
    solution, inputs, closest_match, error_message = rescale_generic_grid(input_args, database_path)

    # Build file out
    tab = at.Table(data=[solution['wv'], solution['spectra']])
    fh = StringIO()
    tab.write(fh, format='ascii.no_header')

    if write_table:
        tab.write('generic.dat')

    # Plot
    fig = figure(title='Rescaled Generic Grid Transmission Spectra'.upper(), plot_width=1100, plot_height=400)
    fig.x_range.start = 0.3
    fig.x_range.end = 5
    fig.line(solution['wv'], solution['spectra'], color='Black', line_width=1)
    fig.xaxis.axis_label = 'Wavelength (um)'
    fig.yaxis.axis_label = 'Transit Depth (Rp/R*)^2'

    if write_plot:
        output_file('generic.html')
        save(fig)

    return fig, fh, closest_match, error_message

def rescale_generic_grid(input_args, database_path):
    """ Pulls a model from the generic grid, rescales it,
    and returns the model and wavelength.

    Parameters
    ----------
    input_args : dict
        A dictionary of the form output from the generic grid form.
        If manual input must include :
        r_star : The radius of the star.
        r_planet : The radius of the planet.
        gravity : The gravity.
        temperature : The temperature.
        condensation : local or rainout
        metallicity
        c_o : carbon/oxygen ratio
        haze
        cloud
    database_path : str
        Path to the generic grid database.

    Returns
    -------
    wv : np.array
        Array of wavelength bins.
    spectra : np.array
        Array of the planetary model spectrum.
    inputs : dict
        The dictionary of inputs given to the function.
    closest_match : dict
        A dictionary with the parameters/model name of the closest
        match in the grid.
    error_message : bool, str
        Either False, for no error, or a message about what went wrong.
    """
    error_message = ''
    try:
        # Parameter validation
        # Set up some nasty tuples first
        scaling_space = [('r_star', [0.05, 10000]),
                         ('r_planet', [0.0,  10000]),
                         ('gravity', [.5, 50]),
                         ('temperature', [400, 2600])]

        inputs = {}
        # First check the scaling
        for tup in scaling_space:
            key, space = tup
            val = float(input_args[key])
            if val >= space[0] and val <= space[1]:
                inputs[key] = val
            else:
                error_message = 'One of the scaling parameters was out of range: {}.'.format(key)
                break

        # Map to nearest model key
        temp_range = np.arange(600, 2700, 100)
        grav_range = np.array([5, 10, 20, 50])
        sort_temp = (np.abs(inputs['temperature'] - temp_range)).argmin()
        sort_grav = (np.abs(inputs['gravity'] - grav_range)).argmin()
        model_temp = temp_range[sort_temp]
        input_args['model_temperature'] = '0{}'.format(model_temp)[-4:]
        model_grav = grav_range[sort_grav]
        input_args['model_gravity'] = '0{}'.format(model_grav)[-2:]

        # Check the model parameters
        str_temp_range = ['0400'] + ['0{}'.format(elem)[-4:] for elem in temp_range]
        model_space = [('condensation', ['local', 'rainout']),
                        ('RCF', ['0.25', '0.50', '0.75', '1.00']),
                       ('model_temperature', str_temp_range),
                       ('model_gravity', ['05', '10', '20', '50']),
                       ('metallicity', ['+0.0', '+1.0', '+1.7', '+2.0', '+2.3']),
                       ('c_o', ['0.35', '0.56', '0.70', '1.00']),
                       ('haze', ['0001', '0010', '0150', '1100']),
                       ('cloud', ['0.00', '0.06', '0.20','1.00'])]

        model_key = ''
        for tup in model_space:
            key, space = tup
            if input_args[key] in space:
                inputs[key] = input_args[key]
                model_key += '{}_'.format(inputs[key])
            else:
                error_message = 'One of the model parameters was out of range.'
                break
        model_key = model_key[:-1]

        # Define constants
        boltzmann = 1.380658E-16 # gm*cm^2/s^2 * Kelvin
        permitivity = 1.6726E-24 * 2.3 #g  cgs  Hydrogen + Helium Atmosphere
        optical_depth = 0.56
        r_sun = 69580000000 # cm
        r_jupiter = 6991100000 # cm

        closest_match = {'model_key': model_key, 'model_gravity': model_grav,
                         'model_temperature': model_temp}

        with h5py.File(database_path, 'r') as f:
            # Can't use the final NaN value
            model_wv = f['/wavelength'][...][:-1]
            model_spectra = f['/spectra/{}'.format(model_key)][...][:-1]

        radius_ratio = np.sqrt(model_spectra) * inputs['r_planet']/inputs['r_star']
        r_star = inputs['r_star'] * r_sun
        r_planet = inputs['r_planet'] * r_jupiter
        model_grav = model_grav * 1e2
        inputs['gravity'] = inputs['gravity'] * 1e2

        # Start with baseline based on model parameters
        scale_height = (boltzmann * model_temp) / (permitivity * model_grav)
        r_planet_base = np.sqrt(radius_ratio) * r_sun
        altitude = r_planet_base - (np.sqrt(radius_ratio[2000])*r_sun)
        opacity = optical_depth * np.sqrt((boltzmann * model_temp * permitivity * model_grav) / \
                                          (2 * np.pi * r_planet_base)) * \
                                  np.exp(altitude / scale_height)
        # Now rescale from baseline
        solution = {}
        solution['scale_height'] = (boltzmann * inputs['temperature']) / (permitivity * inputs['gravity'])
        solution['altitude'] = solution['scale_height'] * \
                               np.log10(opacity/optical_depth * \
                                        np.sqrt((2 * np.pi * r_planet) / \
                                                (boltzmann * inputs['temperature'] * inputs['gravity'])))
        solution['radius'] = solution['altitude'] + r_planet

        # Sort data
        sort = np.argsort(model_wv)
        solution['wv'] = model_wv[sort]
        solution['radius'] = solution['radius'][sort]
        solution['spectra'] = (solution['radius']/r_star)**2

    except (KeyError, ValueError) as e:
        error_message = 'One of the parameters to make up the model was missing or out of range.'
        model_key = 'rainout_0400_50_+0.0_0.70_0010_1.00'
        solution = {}
        with h5py.File(database_path) as f:
            solution['wv'] = f['/wavelength'][...][:-1]
            solution['spectra'] = f['/spectra/{}'.format(model_key)][...][:-1]
        closest_match = {'model_key': model_key, 'model_temperature': 400,
                'model_gravity': 50}
        inputs = input_args

    return solution, inputs, closest_match, error_message


import numpy as np
from astropy.io import fits
from astropy.table import Table, vstack
from astropy.wcs import WCS
import os
import argparse
import logging, traceback
import time
import pandas as pd

from bkg_rate_estimation import rate_obj_from_sqltab
from sqlite_funcs import get_conn, write_result, write_results,\
                        timeID2time_dur, write_results_fromSigImg,\
                        update_square_stat, write_square_res_line,\
                        write_square_results
from dbread_funcs import get_rate_fits_tab, guess_dbfname, get_twinds_tab,\
                    get_seeds_tab, get_info_tab, get_files_tab,\
                    get_square_tab, get_full_sqlite_table_as_df
from config import EBINS0, EBINS1, solid_angle_dpi_fname, fp_dir
from flux_models import Plaw_Flux
from minimizers import NLLH_ScipyMinimize_Wjacob, imxy_grid_miner, NLLH_ScipyMinimize
from drm_funcs import DRMs
from ray_trace_funcs import RayTraces, FootPrints
from LLH import LLH_webins
from models import Bkg_Model_wSA, Point_Source_Model, Point_Source_Model_Wuncoded,\
            CompoundModel, Bkg_Model_wFlatA, Point_Source_Model_Binned_Rates
from do_intllh_scan import kum_mode, kum_pdf, kum_logpdf, kum_deriv_logpdf, deriv2_kum_logpdf

# need to read rate fits from DB
# and read twinds
# and read/get event, dmask, and ebins
# then get bkg_llh_obj and a minimizer
# then loop over all time windows
# minimizing nllh and recording bf params

def cli():
    parser = argparse.ArgumentParser()
    parser.add_argument('--evfname', type=str,\
            help="Event data file",
            default=None)
    parser.add_argument('--dmask', type=str,\
            help="Detmask fname",
            default=None)
    parser.add_argument('--job_id', type=int,\
            help="ID to tell it what seeds to do",\
            default=-1)
    parser.add_argument('--Njobs', type=int,\
            help="Total number of jobs submitted",\
            default=64)
    parser.add_argument('--dbfname', type=str,\
            help="Name to save the database to",\
            default=None)
    parser.add_argument('--rt_dir', type=str,\
            help="Directory with ray traces",\
            default=None)
    parser.add_argument('--pcfname', type=str,\
            help="partial coding file name",\
            default='pc_2.img')
    parser.add_argument('--job_fname', type=str,\
            help="File name for table with what imx/y square for each job",\
            default='job_table2.csv')
    parser.add_argument('--rate_fname', type=str,\
            help="Rate results file name",\
            default='rate_seeds2.csv')
    parser.add_argument('--bkg_fname', type=str,\
            help="Name of the file with the bkg fits",\
            default='bkg_estimation.csv')
    parser.add_argument('--pix_fname', type=str,\
            help="Name of the file with good imx/y coordinates",\
            default='good_pix2scan.npy')
    parser.add_argument('--sim_dir', type=str,\
            help="Name of the simulation directory",\
            default=None)
    parser.add_argument('--log_fname', type=str,\
            help="Name for the log file",\
            default='llh_sim')
    parser.add_argument('--dur_min', type=float,\
            help="Min duration to use",
            default=0.5)
    parser.add_argument('--dur_max', type=float,\
            help="Max duration to use",
            default=4.096)
    parser.add_argument('--dt_min', type=float,\
            help="Min time from trig time to use",
            default=-6.144)
    parser.add_argument('--dt_max', type=float,\
            help="Max time from trig time to use",
            default=4.096)
    args = parser.parse_args()
    return args

def parse_bkg_csv(bkg_fname, solid_angle_dpi, ebins0, ebins1, bl_dmask, rt_dir):

    bkg_df = pd.read_csv(bkg_fname)
    col_names = bkg_df.columns
    nebins = len(ebins0)

    PSnames = []
    for name in col_names:
        if '_imx' in name:
            PSnames.append(name.split('_')[0])
    print PSnames
    Nsrcs = len(PSnames)
    if Nsrcs > 0:
        bkg_name = 'Background_'
    else:
        bkg_name = ''

    bkg_mod = Bkg_Model_wFlatA(bl_dmask, solid_angle_dpi, nebins, use_deriv=True)

    ps_mods = []

    if Nsrcs > 0:
        rt_obj = RayTraces(rt_dir)
        for i in range(Nsrcs):
            name = PSnames[i]
            imx = bkg_df[name+'_imx'][0]
            imy = bkg_df[name+'_imy'][0]
            mod = Point_Source_Model_Binned_Rates(imx, imy, 0.1,\
                                                  [ebins0,ebins1], rt_obj, bl_dmask,\
                                                  use_deriv=True, name=name)
            ps_mods.append(mod)

    return bkg_df, bkg_name, PSnames, bkg_mod, ps_mods

def do_analysis(sim_params_df, sim_tab, twind_df,\
                pl_flux, drm_obj, rt_dir, fp_dir,\
                ev_data, bl_dmask, ebins0, ebins1,\
                conn, db_fname, trigger_time,\
                work_dir, sim_dir,\
                bkg_fname, TSwrite=4.5):

    conn.close()

    nebins = len(ebins0)

    solid_ang_dpi = np.load(solid_angle_dpi_fname)

    bkg_miner = NLLH_ScipyMinimize('')
    sig_miner = NLLH_ScipyMinimize_Wjacob('')

    bkg_df, bkg_name, PSnames, bkg_mod, ps_mods =\
            parse_bkg_csv(bkg_fname, solid_ang_dpi,\
                        ebins0, ebins1, bl_dmask, rt_dir)

    bkg_mod.has_deriv = False
    bkg_mod_list = [bkg_mod]
    Nsrcs = len(ps_mods)
    if Nsrcs > 0:
        bkg_mod_list += ps_mods
        for ps_mod in ps_mods:
            ps_mod.has_deriv = False
        bkg_mod = CompoundModel(bkg_mod_list)

    for sim_param_ind, sim_param_row in sim_params_df.iterrows():

        logging.info("Starting params_id: %d"%(sim_param_row['params_id']))

        bl = (sim_tab['params_id']==sim_param_row['params_id'])
        sim_df = sim_tab[bl]

        rt_obj = RayTraces(rt_dir, max_nbytes=6e9)
        fp_obj = FootPrints(fp_dir)

        res_dfs2write = []

        for sim_ind, sim_row in sim_df.iterrows():

            logging.debug("Starting simid: %d, table index: %d"\
                        %(sim_row['simid'],sim_ind))
            tstart = sim_row['tstart']
            tstop = tstart + sim_row['dur']
            dur = sim_row['dur']

            imx_sim = sim_row['imx']
            imy_sim = sim_row['imy']

            ev_sim_tab = Table.read(sim_row['fname'])

            evdata_ = ev_data.copy()
            evdata = vstack([evdata_, ev_sim_tab])
            evdata.sort('TIME')

            bkg_llh_obj = LLH_webins(evdata, ebins0, ebins1, bl_dmask)
            sig_llh_obj = LLH_webins(evdata, ebins0, ebins1, bl_dmask)

            xax = np.linspace(-1e-3, 1e-3, 2+1) + imx_sim
            yax = np.linspace(-1e-3, 1e-3, 2+1) + imy_sim
            grids = np.meshgrid(xax, yax)
            imxs = grids[0].ravel()
            imys = grids[1].ravel()
            Npix = len(imxs)

            bl = (twind_df['time']<(tstop-.1*dur))&\
            ((twind_df['duration']+twind_df['time'])>(tstart+.1*dur))
            tgrps = twind_df[bl].groupby('timeID')
            logging.debug("%d timeIDs to do"%(len(tgrps)))

            res_dicts = []

            for timeID, tdf in tgrps:

                res_dict = {}
                res_dict['params_id'] = sim_param_row['params_id']
                res_dict['simid'] = sim_row['simid']
                res_dict['timeID'] = timeID

                t0 = tdf['time'].values[0]
                dt = tdf['duration'].values[0]
                tmid = t0 + dt/2.
                t1 = t0 + dt
                res_dict['time'] = t0
                res_dict['duration'] = dt

                bkg_llh_obj.set_time(t0, t1)
                sig_llh_obj.set_time(t0, t1)

                bkg_row = bkg_df.iloc[np.argmin(np.abs(tmid - bkg_df['time']))]

                bkg_llh_obj.set_model(bkg_mod)

                bkg_miner.set_llh(bkg_llh_obj)

                bkg_params = {pname:bkg_row[pname] for pname in\
                            bkg_llh_obj.model.param_names}
                # bkg_miner.set_fixed_params(bkg_llh_obj.model.param_names)
                bkg_miner.set_fixed_params(bkg_params.keys(), values=bkg_params.values())

                # bkg_params = {pname:bkg_llh_obj.model.param_dict[pname]['val'] for\
                #                 pname in bkg_llh_obj.model.param_names}
                bkg_nllh = -bkg_llh_obj.get_llh(bkg_params)
                res_dict['bkg_nllh'] = bkg_nllh
                logging.debug("bkg_param_dict: ")
                logging.debug(bkg_miner.param_info_dict)
                logging.debug("bkg_nllh: %.3f" %(bkg_nllh))
                imx_, imy_ = np.nanmean(imxs), np.nanmean(imys)

                # sig_mod = Point_Source_Model(imx_,\
                #                         imy_, 0.3,\
                #                         pl_flux, drm_obj,\
                #                         [ebins0,ebins1], rt_obj, bl_dmask,\
                #                         use_deriv=True)

                sig_mod = Point_Source_Model_Wuncoded(imx_, imy_, 0.3,\
                                    pl_flux, drm_obj, [ebins0,ebins1],\
                                    rt_obj, fp_obj, bl_dmask, use_deriv=True)

                sig_mod.drm_im_update = .2
                comp_mod = CompoundModel([bkg_mod, sig_mod])
                sig_llh_obj.set_model(comp_mod)
                sig_miner.set_llh(sig_llh_obj)
                fixed_pars = [pname for pname in sig_miner.param_names if\
                            ('A' not in pname) or ('gamma' not in pname)]
                sig_miner.set_fixed_params(fixed_pars)
                sig_miner.set_fixed_params(['Signal_A', 'Signal_gamma'], fixed=False)

                TSs = np.zeros(Npix)
                sig_nllhs = np.zeros(Npix)
                As = np.zeros(Npix)
                gammas = np.zeros(Npix)

                for ii in range(Npix):
                    try:
                        sig_miner.set_fixed_params(['Signal_imx', 'Signal_imy'], [imxs[ii],imys[ii]])
                        pars, nllh, res = sig_miner.minimize()
                        TS = np.sqrt(2.*(bkg_nllh - nllh[0]))
                        # if TS >= TS_min:
                        if np.isnan(TS):
                            TS = 0.0
                        TSs[ii] = TS
                        sig_nllhs[ii] = nllh[0]
                        As[ii] = pars[0][0]
                        gammas[ii] = pars[0][1]

                    except Exception as E:
                        logging.error(E)
                        logging.error(traceback.format_exc())
                        logging.error("Failed to minimize seed: ")
                        logging.error((imxs[ii],imys[ii]))

                logging.debug("Max TS: %.2f" %(np.nanmax(TSs)))

                best_ind = np.nanargmax(TSs)
                res_dict['TS'] = TSs[best_ind]
                res_dict['imx'] = imxs[best_ind]
                res_dict['imy'] = imys[best_ind]
                res_dict['A'] = As[best_ind]
                res_dict['ind'] = gammas[best_ind]
                res_dict['sig_nllh'] = sig_nllhs[best_ind]
                res_dicts.append(res_dict)
                # fname = os.path.join(work_dir,\
                #         'res_%d_%d_.fits' %(res_dict['timeID'],\
                #         res_dict['squareID']))

                # TSbl = (TSs>=TSwrite)
                # if np.sum(TSbl) > 0:
                #     logging.info("%d above TS of %.1f"%(np.sum(TSbl),TSwrite))
                #     res_dict['TS'] = TSs[TSbl]
                #     res_dict['imx'] = imxs[TSbl]
                #     res_dict['imy'] = imys[TSbl]
                #     res_dict['A'] = As[TSbl]
                #     res_dict['ind'] = gammas[TSbl]
                #     res_dict['sig_nllh'] = sig_nllhs[TSbl]
                #     # res_dict['fname'] = fname
                #     res_dfs2write.append(pd.DataFrame(res_dict))

            res_df = pd.DataFrame(res_dicts)
            res_dfs2write.append(res_df)

        fname = os.path.join(sim_dir,\
                'res_paramsID_%d_.csv' %(sim_param_row['params_id']))

        res_df = pd.concat(res_dfs2write, ignore_index=True)
        res_df.to_csv(fname, index=False)
        logging.info("Saved results to")
        logging.info(fname)

def main(args):

    # fname = 'llh_analysis_from_rate_seeds_' + str(args.job_id)
    fname = args.log_fname + '_' + str(args.job_id)

    logging.basicConfig(filename=fname+'.log', level=logging.DEBUG,\
                    format='%(asctime)s-' '%(levelname)s- %(message)s')

    t_0 = time.time()

    if args.dbfname is None:
        db_fname = guess_dbfname()
        if isinstance(db_fname, list):
            db_fname = db_fname[0]
    else:
        db_fname = args.dbfname

    logging.info('Connecting to DB')
    conn = get_conn(db_fname)

    info_tab = get_info_tab(conn)
    logging.info('Got info table')

    files_tab = get_files_tab(conn)
    logging.info('Got files table')

    trigtime = info_tab['trigtimeMET'][0]

    evfname = files_tab['evfname'][0]
    # ev_data = fits.open(evfname)[1].data
    ev_data = Table.read(evfname)
    dmask_fname = files_tab['detmask'][0]
    dmask = fits.open(dmask_fname)[0].data
    bl_dmask = (dmask==0.0)
    logging.debug('Opened up event and detmask files')

    bkg_fits_df = pd.read_csv(args.bkg_fname)

    sim_params_df = pd.read_csv(os.path.join(args.sim_dir, 'sim_param_table.csv'))
    sim_tab = pd.read_csv(os.path.join(args.sim_dir, 'sim_table.csv'))

    bl = (sim_params_df['im_id']==args.job_id)
    sim_params_df = sim_params_df[bl]
    Nsims = len(sim_params_df)

    # rate_fits_df = get_rate_fits_tab(conn)
    # bkg_rates_obj = rate_obj_from_sqltab(rate_fits_df, 0, 1)

    time_starting = time.time()
    proc_num = args.job_id
    # init classes up here

    drm_dir = files_tab['drmDir'][0]
    if args.rt_dir is None:
        rt_dir = files_tab['rtDir'][0]
    else:
        rt_dir = args.rt_dir
    drm_obj = DRMs(drm_dir)
    # rt_obj = RayTraces(rt_dir, max_nbytes=1e10)
    work_dir = files_tab['workDir'][0]

    pl_flux = Plaw_Flux()

    ebins0 = np.array(EBINS0)
    ebins1 = np.array(EBINS1)
    logging.debug("ebins0")
    logging.debug(ebins0)
    logging.debug("ebins1")
    logging.debug(ebins1)

    twind_df = get_twinds_tab(conn)

    bl = ((twind_df['time']-trigtime)>=args.dt_min)&\
    ((twind_df['duration']+twind_df['time']-trigtime)<args.dt_max)&\
    (twind_df['duration']>=args.dur_min)&(twind_df['duration']<=args.dur_max)

    twind_df = twind_df[bl]
    logging.info("Got TimeWindows table")

    do_analysis(sim_params_df, sim_tab, twind_df, pl_flux,\
                    drm_obj, rt_dir, fp_dir,\
                    ev_data, bl_dmask, ebins0, ebins1,\
                    conn, db_fname, trigtime,\
                    work_dir, args.sim_dir, args.bkg_fname)
    # do_analysis(square_tab, rate_res_tab, good_pix['imx'], good_pix['imy'], pl_flux,\
    #                 drm_obj, rt_dir,\
    #                 bkg_llh_obj, sig_llh_obj,\
    #                 conn, db_fname, trigtime, work_dir,bkg_fits_df)
    conn.close()

if __name__ == "__main__":

    args = cli()

    main(args)