#!/usr/bin/env python3
import sys
import os
import json
import networkx as nx
import math
from networkx.algorithms import approximation as approx
import time
from itertools import groupby as g
from operator import itemgetter
from graphFunctions import *
import csv
import configparser
import logging
#from networkx.readwrite import d3_js
from networkx.readwrite import json_graph
#import mpld3
#mpld3.enable_notebook()
#from mpld3 import plugins
#import matplotlib.pyplot as plt

Config = configparser.ConfigParser()
Config.read("../etc/config.cfg")

logging.basicConfig(filename=Config.get("Debug", "Logfile"),level=logging.INFO, format='%(asctime)s %(message)s')

debug = Config.getboolean("Debug", "DebugInfo")
standardDirectory = Config.get("Path", "StandardDirectory")
filename = ""
lastcomm = ""
outputfile = ""
maxEdgeValue = Config.getint("Graph", "maxEdgeValue");
# Values for Edge creation
createEdgePercent = Config.getint("Graph", "createEdgePercent") 			# Bis zu einem Wert von 30% NICHT ähnlich
createBlueEdgePercent = Config.getint("Graph", "createBlueEdgePercent") 	# Bis zu einem Wert von 50% ETWAS ähnlich
createBluedEdges = True

for argument in sys.argv:
	if lastcomm.strip() == "-f":
		filename = argument.strip();
	if lastcomm.strip() == "-d":
		debug = True
	if lastcomm.strip() == "-m":
		maxEdgeValue = int(argument.strip());
	if lastcomm.strip() == "-o":
		outputfile = argument.strip();
	lastcomm = argument
if lastcomm.strip() == "-f":
	filename = argument.strip();
if lastcomm.strip() == "-d":
	debug = True
if lastcomm.strip() == "-m":
	maxEdgeValue = int(argument.strip());
if lastcomm.strip() == "-o":
	outputfile = argument.strip();

logging.debug (" ========")
logging.info ("  Starting PS-Clustering Algorithm with Complement on "+filename)
logging.info (" ========")
start_time = time.time()

if filename == "":
	logging.error ("No file name given.")
	exit()

G=nx.Graph()
logging.info (" ... now reading graph at "+filename)
G = nx.read_gml(filename)

if len(G) == 0:
	logging.error (" ... Graph is empty! Exiting!")
	exit()
build_time = time.time() - start_time
logging.info (" => Build time: "+str(build_time))

start_time = time.time()

G2 = G.copy()
G2 = nx.complement(G2)
nh = getNeighbors(G2)
cliques = []
for key, value in nh.items():
	u1 = key
	#print (u1)
	if u1 not in G2.nodes():
		#print (" -> not in G2")
		continue
	u2 = getSmalestNeighbor(G2,u1)
	if u2 == None:
		#print (" -> None ")
		continue
	G2.add_node (str(u1)+"|"+str(u2))
	for node in G2[u1]:
		if nodes_connected(G2, u2, node):
			G2.add_edge(str(u1)+"|"+str(u2), node)
	if 'clique' in G2.node[u2]:
		G2.node[str(u1)+"|"+str(u2)]['clique']=G2.node[u2]['clique']
	else:
		G2.node[str(u1)+"|"+str(u2)]['clique']=getMaxClique(G2)+1
	G2.remove_node(u1)
	G2.remove_node(u2)
for node in G2.nodes():
	if 'clique' not in G2.node[node]:
		G2.node[node]['clique']=getMaxClique(G2)+1
cl = nx.get_node_attributes(G2,'clique')
for key, value in cl.items():
	if "|" in str(key):
		nodes = key.split("|")
	else:
		nodes = [key]
	for node in nodes:
		if node in G.nodes():
			G.node[node]['color']=value
		elif int(node) in G.nodes():
				G.node[int(node)]['color']=value
		else:
			print (" Error, node does not exist! ")

#d = nx.coloring.greedy_color(G, strategy=nx.coloring.strategy_largest_first)
d = nx.get_node_attributes(G, 'color')
colorCount = d[max(d, key=lambda key: d[key])]

build_time = time.time() - start_time
logging.info (" => Clique time: "+str(build_time))
start_time = time.time()

start_time = time.time()
# valueGraph und Farbenliste bauen
valueGraph = buildValueGraph (G, d )
outputpathfilename = standardDirectory+outputfile
if os.path.isabs (outputfile):
	outputpathfilename = outputfile
logging.debug ("  ... saving weightes Value Graph at "+outputpathfilename+'coloredWithValue.gml')
nx.write_gml(valueGraph, outputpathfilename+'coloredWithValue.gml')
# Build a colorlist and sort min first
colorList = []
for i in range(0,colorCount):
		countC = countColor (G,i);
		colorList.append ( [i, countC] )
colorList = sorted(colorList,key=itemgetter(1))

build_time = time.time() - start_time
logging.info (" => Build weighted graph time: "+str(build_time))
start_time = time.time()

#G = addBlueEdges (G, createEdgePercent, createBlueEdgePercent, createBluedEdges, data)

# For all nodes
countNodesEliminated = 0;
listConnections = []

colorCount = d[max(d, key=lambda key: d[key])]
minYear = 3000
maxYear = 0
for i in range(0,colorCount):#
	yearc = getClusterYearsList (G,i)
	for year in yearc:#
		if year[0]==0:
			continue
		if year[0]>maxYear:
			maxYear = year[0]
		if year[0]<minYear:
			minYear = year[0]

for node in G.nodes():
	G.node[node]['year']=int(G.node[node]['year'])

for colorSet in colorList:
	# Get a list of nodes with that color
	tmpNodeList = list(n for n,d in G.nodes_iter(data=True) if d['color']==colorSet[0])
	#print (" Color count before: "+str(countColor(G, colorSet[0])))
	# Now iterate over these nodes
	for i in range(0,len(tmpNodeList)):
		# catch that node
		node = tmpNodeList[i]
		# check if this node is not already an end-point!
		if G.node[node]["end"]==0:
			colors = [];
			# for all neighbours
			for nb in G.neighbors(node):
				# is edge blue?
				if G[node][nb]["color"] == "blue":
					# is node not yet excluded from stable sets?
					if (G.node[nb]["color"] >= 0):
						# Nodes are not in the same stable sets
						if (G.node[nb]["color"] != G.node[node]["color"]):
							colors.append (G.node[nb]["color"])
			# Now remove all double entries
			uniqueList = unique(colors)
			# check if "node" has at last two blue neighbours in different stable sets
			if len(uniqueList) >= 2:
				# Node is excluded from stable sets (color -1)
				G.node[node]["color"] = -1;
				countNodesEliminated += 1
				# now we see, that this node connects all stable sets connected with blue edges:
				for i in range(0,len(uniqueList)) :
					for j in range (i+1, len(uniqueList)):
						newYear = "y"+str(G.node[node]["year"])
						#print (" Year: "+newYear + " from "+ str(G.node[node]["year"]))
						if str(uniqueList[i]-1) in valueGraph.nodes() and str(str(uniqueList[j]-1)) in valueGraph.nodes():
							if valueGraph.has_edge(str(uniqueList[i]-1),str(uniqueList[j]-1)):
								valueGraph[str(uniqueList[i]-1)][str(uniqueList[j]-1)]["value"]+=1
								if  "y"+str(G.node[node]["year"]) in valueGraph[str(uniqueList[i]-1)][str(uniqueList[j]-1)]:
									# Jahr existiert
									valueGraph[str(uniqueList[i]-1)][str(uniqueList[j]-1)]["y"+str(G.node[node]["year"])] += 1
								else:
									valueGraph.add_edge (str(uniqueList[i]-1) , str(uniqueList[j]-1), {newYear : 1})
							else:
								valueGraph.add_edge(str(uniqueList[i]-1),str(uniqueList[j]-1), value=1)
								valueGraph.add_edge (str(uniqueList[i]-1) , str(uniqueList[j]-1), {newYear : 1})
							for yearsum in range(minYear, G.node[node]["year"]-1):
								if  "s"+str(yearsum) not in valueGraph[str(uniqueList[i]-1)][str(uniqueList[j]-1)]:
									newYear2 = "s"+str(yearsum)
									valueGraph.add_edge (str(uniqueList[i]-1) , str(uniqueList[j]-1), {newYear2 : 0})
							for yearsum in range(G.node[node]["year"], maxYear):
								if  "s"+str(yearsum) in valueGraph[str(uniqueList[i]-1)][str(uniqueList[j]-1)]:
									valueGraph[str(uniqueList[i]-1)][str(uniqueList[j]-1)]["s"+str(yearsum)] += 1
								else:
									newYear2 = "s"+str(yearsum)
									valueGraph.add_edge (str(uniqueList[i]-1) , str(uniqueList[j]-1), {newYear2 : 1})
						#else:
						#	logging.warn (" Error, nodes do not exist. ")
						#valueGraph[str(uniqueList[i]-1)][str(uniqueList[j]-1)]["y"+str(G.node[node]["year"])] += 1
				# mark all nodes as end-nodes
				for nb in G.neighbors(node):
					# is edge blue?
					if G[node][nb]["color"] == "blue":
						# is node not yet excluded from stable sets?
						if (G.node[nb]["color"] >= 0):
							G.node[nb]["end"] = 1
	#print (" Color count after: "+str(countColor(G, colorSet[0])))
#print (str(valueGraph.number_of_nodes()))

if "-1" in valueGraph.nodes():
	valueGraph.remove_node("-1")
if -1 in valueGraph.nodes():
	valueGraph.remove_node(-1)

count = 15

for node in valueGraph.nodes():
	newValue = countColor (G, int(node))
	if newValue == 0:
		#print ("Removing node "+str(node)+" with "+str(valueGraph.node[node]["text"]))
		valueGraph.remove_node(node)
	else:
		valueGraph.node[node]["value"]=newValue
		terms = getClusterName(G,node,count);
		journals = getClusterJournals (G,node,count)
		yearc = getClusterYearsList (G,node,count)
		valueGraph.node[node]["text"]=terms
		valueGraph.node[node]["journal"]=journals
		yearc.sort(key=lambda tup: tup[0])
		#print (str(yearc))
		summe = 0
		#print ("New Sum")
		#print (str(minYear) +" -- "+str(maxYear))
		#print (yearc)
		for year in range(minYear, maxYear):
			addValue = 0
			#print ("Looking for year "+str(year))
			for years in yearc:
				if years[0]==0:
					continue
				if years[0] == year:
					#print (str(years[0]) +" == "+str(year))
					addValue = years[1]
			#print (" "+str(year))
			valueGraph.node[node]["y"+str(year)] = addValue
			valueGraph.node[node]["s"+str(year)] = addValue + summe
			summe += addValue
		#csvwriter.writerow ([''])
		#csvwriter.writerow (["Node", int(node), "Nodes:", newValue])
		#csvwriter.writerow (getClusterName (G, int(node), maxCount=-1, separator=";").split(";"))
		#csvwriter.writerow ([''])
		#nodeList = (n for n in G if G.node[n]['color']==int(node))
		#for nodeG in nodeList:
		#	csvwriter.writerow ([ G.node[nodeG]['title'], G.node[nodeG]['ref'], G.node[nodeG]['uri']])
logging.debug ("  ... saving PS Value Graph at "+outputpathfilename+'coloredWithValuePS.gml')
nx.write_gml(valueGraph, outputpathfilename+'coloredWithValuePS.gml')

build_time = time.time() - start_time
logging.debug (" => PS time: "+str(build_time))
start_time = time.time()

ps_time = build_time

#csvwriter.writerows (sorted(oldList.items(), key=operator.itemgetter(1), reverse=True))

logging.debug ("  ... saving Network RB at  "+outputpathfilename+'network_rb.gml')
# Output nach neu färben
nx.write_gml(G, outputpathfilename+'network_rb.gml')

numbers = getMeshTermNumber (G);
#print ( sorted(numbers.items(), key=operator.itemgetter(1)))

# Export
data = json_graph.node_link_data(valueGraph)
with open(outputpathfilename+'coloredWithValuePS.json', 'w') as outfile:
    json.dump(data, outfile)

#fig, axs = plt.subplots(1, 1, figsize=(10, 10))
#ax = axs
#pos = None
#mpld3.plugins.connect(fig,  NetworkXD3ForceLayout(valueGraph,
#                                                  pos,
#                                                  ax,
#                                                  gravity=.5,
#                                                  link_distance=20,
#                                                  charge=-600,
#                                                  friction=1
#                                                 )
#                     )

if debug:
	#logging.debug ("Graph created. Number of nodes: " + str(G.number_of_nodes())+ ", number edges: "+ str(G.number_of_edges()))
	#logging.debug (" Blue edges: "+str(bluedEdges))
	#print (" Black edges: "+ str(countBlack) + ", Blued edges: "+ str(countBlue) +", lost edged: "+str(countLost))
	#logging.debug (" Build time: \t " + str((build_time)))
	#logging.debug (" Color time: \t " + str((color_time )))
	logging.info (" PS time:  \t " + str((ps_time )))

	logging.info (" ========")
	logging.info (" Farben: \t\t "+str(colorCount))
	logging.info (" minMPS: \t\t "+str(valueGraph.number_of_nodes()-1))
	#lower = colorCount - math.floor (bluedEdges / 2)
	#if lower < 2:
	#	lower = 2
	#logging.debug (" l_b: \t \t\t "+str(lower))
	#print (" Cliquenzahl:  "+ str(clique) )
	logging.info (" Eliminated nodes: \t "+str(countNodesEliminated))
	#print (" Clique time: " + str((clique_time - color_time)))
	#print (G.edges())
#nx.set_node_attributes(G, 'label', {0: "0", 1: "1"})

output = []
output.append (ps_time)
output.append (colorCount)
output.append (valueGraph.number_of_nodes())
output.append (countNodesEliminated)

print (output)


# cython: language_level=3
# distutils: language = c++
# -*- coding: utf-8 -*-
# *****************************************************************************
# Copyright (c) 2016-2020, Intel Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# - Redistributions of source code must retain the above copyright notice,
#   this list of conditions and the following disclaimer.
# - Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
# *****************************************************************************

"""
Interface of the Linear Algebra part of the DPNP

Notes
-----
This module is a face or public interface file for the library
it contains:
 - Interface functions
 - documentation for the functions
 - The functions parameters check

"""

import dpnp
import numpy

from dpnp.dparray import dparray
from dpnp.dpnp_utils import *
from dpnp.linalg.dpnp_algo_linalg import *

__all__ = [
    "cholesky",
    "cond",
    "det",
    "eig",
    "eigvals",
    "inv",
    "matrix_power",
    "matrix_rank",
    "multi_dot",
    "norm",
    "qr",
    "svd",
]

def cholesky(input):
    """
    Cholesky decomposition.
    Return the Cholesky decomposition, `L * L.H`, of the square matrix `input`,
    where `L` is lower-triangular and .H is the conjugate transpose operator
    (which is the ordinary transpose if `input` is real-valued).  `input` must be
    Hermitian (symmetric if real-valued) and positive-definite. No
    checking is performed to verify whether `a` is Hermitian or not.
    In addition, only the lower-triangular and diagonal elements of `input`
    are used. Only `L` is actually returned.

    Parameters
    ----------
    input : (..., M, M) array_like
        Hermitian (symmetric if all elements are real), positive-definite
        input matrix.

    Returns
    -------
    L : (..., M, M) array_like
        Upper or lower-triangular Cholesky factor of `input`.  Returns a
        matrix object if `input` is a matrix object.
    """

    if not use_origin_backend(input):
        if not isinstance(input, dparray):
            pass
        elif input.shape[-1] != input.shape[-2]:
            pass
        elif input.ndim < 3:
            pass
        else:
            return dpnp_cholesky(input)

    return call_origin(numpy.linalg.cholesky, input)

def cond(input, p=None):
    """
    Compute the condition number of a matrix.
    For full documentation refer to :obj:`numpy.linalg.cond`.

    Limitations
    -----------
    Input array is supported as :obj:`dpnp.ndarray`.
    Parameter p=[None, 1, -1, 2, -2, numpy.inf, -numpy.inf, 'fro'] is supported.

    See Also
    --------
    :obj:`dpnp.norm` : Matrix or vector norm.
    """

    is_input_dparray = isinstance(input, dparray)

    if (not use_origin_backend(input) and is_input_dparray):
        if p in [None, 1, -1, 2, -2, numpy.inf, -numpy.inf, 'fro']:
            result = dpnp_cond(input, p=p)
            return result.dtype.type(result[0])
        else:
            pass

    return call_origin(numpy.linalg.cond, input, p)

def det(input):
    """
    Compute the determinant of an array.

    Parameters
    ----------
    input : (..., M, M) array_like
        Input array to compute determinants for.

    Returns
    -------
    det : (...) array_like
        Determinant of `input`.
    """
    is_input_dparray = isinstance(input, dparray)

    if not use_origin_backend(input) and is_input_dparray:
        if input.shape[-1] == input.shape[-2]:
            result = dpnp_det(input)

            # scalar returned
            if result.shape == (1,):
                return result.dtype.type(result[0])

            return result

    return call_origin(numpy.linalg.det, input)

def eig(x1):
    """
    Compute the eigenvalues and right eigenvectors of a square array.

    .. seealso:: :obj:`numpy.linalg.eig`

    """

    is_x1_dparray = isinstance(x1, dparray)

    if (not use_origin_backend(x1) and is_x1_dparray):
        if (x1.size > 0):
            return dpnp_eig(x1)

    return call_origin(numpy.linalg.eig, x1)

def eigvals(input):
    """
    Compute the eigenvalues of a general matrix.
    Main difference between `eigvals` and `eig`: the eigenvectors aren't
    returned.

    Parameters
    ----------
    input : (..., M, M) array_like
        A complex- or real-valued matrix whose eigenvalues will be computed.

    Returns
    -------
    w : (..., M,) ndarray
        The eigenvalues, each repeated according to its multiplicity.
        They are not necessarily ordered, nor are they necessarily
        real for real matrices.
    """

    is_input_dparray = isinstance(input, dparray)

    if (not use_origin_backend(input) and is_input_dparray):
        if (input.size > 0):
            return dpnp_eigvals(input)

    return call_origin(numpy.linalg.eigvals, input)

def inv(input):
    """
    Divide arguments element-wise.

    For full documentation refer to :obj:`numpy.linalg.inv`.

    Limitations
    -----------
        Input array is supported as :obj:`dpnp.ndarray`.
        Dimension of input array is supported to be equal to ``2``.
        Shape of input array is limited by ``input.shape[0] == input.shape[1]``, ``input.shape[0] >= 2``.
        Otherwise the function will be executed sequentially on CPU.
    """

    is_input_dparray = isinstance(input, dparray)

    if (not use_origin_backend(input) and is_input_dparray):
        if input.ndim == 2 and input.shape[0] == input.shape[1] and input.shape[0] >= 2:
            return dpnp_inv(input)

    return call_origin(numpy.linalg.inv, input)

def matrix_power(input, count):
    """
    Raise a square matrix to the (integer) power `count`.

    Parameters
    ----------
    input : sequence of array_like

    Returns
    -------
    output : dparray
        Returns the dot product of the supplied arrays.

    See Also
    --------
    :obj:`numpy.linalg.matrix_power`

    """

    is_input_dparray = isinstance(input, dparray)

    if not use_origin_backend(input) and is_input_dparray and count > 0:
        result = input
        for id in range(count - 1):
            result = dpnp.matmul(result, input)

        return result

    input1 = dpnp.asnumpy(input) if is_input_dparray else input

    # TODO need to put dparray memory into NumPy call
    result_numpy = numpy.linalg.matrix_power(input1, count)
    result = result_numpy
    if isinstance(result, numpy.ndarray):
        result = dparray(result_numpy.shape, dtype=result_numpy.dtype)
        for i in range(result.size):
            result._setitem_scalar(i, result_numpy.item(i))

    return result

def matrix_rank(input, tol=None, hermitian=False):
    """
    Return matrix rank of array
    Rank of the array is the number of singular values of the array that are
    greater than `tol`.

    Parameters
    ----------
    M : {(M,), (..., M, N)} array_like
        Input vector or stack of matrices.
    tol : (...) array_like, float, optional
        Threshold below which SVD values are considered zero. If `tol` is
        None, and ``S`` is an array with singular values for `M`, and
        ``eps`` is the epsilon value for datatype of ``S``, then `tol` is
        set to ``S.max() * max(M.shape) * eps``.
    hermitian : bool, optional
        If True, `M` is assumed to be Hermitian (symmetric if real-valued),
        enabling a more efficient method for finding singular values.
        Defaults to False.

    Returns
    -------
    rank : (...) array_like
        Rank of M.

    """

    is_input_dparray = isinstance(input, dparray)

    if not use_origin_backend(input) and is_input_dparray:
        if tol is not None:
            checker_throw_value_error("matrix_rank", "tol", type(tol), None)
        if hermitian is not False:
            checker_throw_value_error("matrix_rank", "hermitian", hermitian, False)

        result = dpnp_matrix_rank(input)

        # scalar returned
        if result.shape == (1,):
            return result.dtype.type(result[0])

        return result

    return call_origin(numpy.linalg.matrix_rank, input, tol, hermitian)

def multi_dot(arrays, out=None):
    """
    Compute the dot product of two or more arrays in a single function call

    Parameters
    ----------
    arrays : sequence of array_like
        If the first argument is 1-D it is treated as row vector.
        If the last argument is 1-D it is treated as column vector.
        The other arguments must be 2-D.
    out : ndarray, optional
        unsupported

    Returns
    -------
    output : ndarray
        Returns the dot product of the supplied arrays.

    See Also
    --------
    :obj:`numpy.multi_dot`

    """

    n = len(arrays)

    if n < 2:
        checker_throw_value_error("multi_dot", "arrays", n, ">1")

    result = arrays[0]
    for id in range(1, n):
        result = dpnp.dot(result, arrays[id])

    return result

def norm(input, ord=None, axis=None, keepdims=False):
    """
    Matrix or vector norm.
    This function is able to return one of eight different matrix norms,
    or one of an infinite number of vector norms (described below), depending
    on the value of the ``ord`` parameter.

    Parameters
    ----------
    input : array_like
        Input array.  If `axis` is None, `x` must be 1-D or 2-D, unless `ord`
        is None. If both `axis` and `ord` are None, the 2-norm of
        ``x.ravel`` will be returned.
    ord : optional
        Order of the norm (see table under ``Notes``). inf means numpy's
        `inf` object. The default is None.
    axis : optional.
        If `axis` is an integer, it specifies the axis of `x` along which to
        compute the vector norms.  If `axis` is a 2-tuple, it specifies the
        axes that hold 2-D matrices, and the matrix norms of these matrices
        are computed.  If `axis` is None then either a vector norm (when `x`
        is 1-D) or a matrix norm (when `x` is 2-D) is returned. The default
        is None.
    keepdims : bool, optional
        If this is set to True, the axes which are normed over are left in the
        result as dimensions with size one.  With this option the result will
        broadcast correctly against the original `x`.

    Returns
    -------
    n : float or ndarray
        Norm of the matrix or vector(s).
    """

    if not use_origin_backend(input):
        if not isinstance(input, dparray):
            pass
        elif not isinstance(axis, int) and not isinstance(axis, tuple) and axis is not None:
            pass
        elif keepdims is not False:
            pass
        elif ord not in [None, 0, 3, 'fro', 'f']:
            pass
        else:
            result = dpnp_norm(input, ord=ord, axis=axis)

            # scalar returned
            if result.shape == (1,) and axis is None:
                return result.dtype.type(result[0])

            return result

    return call_origin(numpy.linalg.norm, input, ord, axis, keepdims)

#linalg.qr(a, mode='reduced')
def qr(a, mode='complete'):
    """
    Compute the qr factorization of a matrix.

    Factor the matrix `a` as *qr*, where `q` is orthonormal and `r` is
    upper-triangular.

    For full documentation refer to :obj:`numpy.linalg.qr`.

    Limitations
    -----------
    Input array is supported as :obj:`dpnp.ndarray`.
    Parameter mode='complete' is supported.

    """

    if not use_origin_backend(a):
        if not isinstance(a, dparray):
            pass
        elif not mode == 'complete':
            pass
        else:
            return dpnp_qr(a, mode)

    return call_origin(numpy.linalg.qr, a, mode)

def svd(a, full_matrices=True, compute_uv=True, hermitian=False):
    """
    Singular Value Decomposition.

    For full documentation refer to :obj:`numpy.linalg.svd`.

    Examples
    --------
    >>> import dpnp as np
    >>> a = np.random.randn(9, 6) + 1j*np.random.randn(9, 6)
    >>> b = np.random.randn(2, 7, 8, 3) + 1j*np.random.randn(2, 7, 8, 3)

    Reconstruction based on full SVD, 2D case:

    >>> u, s, vh = np.linalg.svd(a, full_matrices=True)
    >>> u.shape, s.shape, vh.shape
    ((9, 9), (6,), (6, 6))
    >>> np.allclose(a, np.dot(u[:, :6] * s, vh))
    True
    >>> smat = np.zeros((9, 6), dtype=complex)
    >>> smat[:6, :6] = np.diag(s)
    >>> np.allclose(a, np.dot(u, np.dot(smat, vh)))
    True

    Reconstruction based on reduced SVD, 2D case:

    >>> u, s, vh = np.linalg.svd(a, full_matrices=False)
    >>> u.shape, s.shape, vh.shape
    ((9, 6), (6,), (6, 6))
    >>> np.allclose(a, np.dot(u * s, vh))
    True
    >>> smat = np.diag(s)
    >>> np.allclose(a, np.dot(u, np.dot(smat, vh)))
    True

    Reconstruction based on full SVD, 4D case:

    >>> u, s, vh = np.linalg.svd(b, full_matrices=True)
    >>> u.shape, s.shape, vh.shape
    ((2, 7, 8, 8), (2, 7, 3), (2, 7, 3, 3))
    >>> np.allclose(b, np.matmul(u[..., :3] * s[..., None, :], vh))
    True
    >>> np.allclose(b, np.matmul(u[..., :3], s[..., None] * vh))
    True

    Reconstruction based on reduced SVD, 4D case:

    >>> u, s, vh = np.linalg.svd(b, full_matrices=False)
    >>> u.shape, s.shape, vh.shape
    ((2, 7, 8, 3), (2, 7, 3), (2, 7, 3, 3))
    >>> np.allclose(b, np.matmul(u * s[..., None, :], vh))
    True
    >>> np.allclose(b, np.matmul(u, s[..., None] * vh))
    True

    """

    if not use_origin_backend(a):
        if not isinstance(a, dparray):
            pass
        elif not a.ndim == 2:
            pass
        elif not full_matrices == True:
            pass
        elif not compute_uv == True:
            pass
        elif not hermitian == False:
            pass
        else:
            return dpnp_svd(a, full_matrices, compute_uv, hermitian)

    return call_origin(numpy.linalg.svd, a, full_matrices, compute_uv, hermitian)


import numpy as np
import array
import os, sys
import re
import time
import multiprocessing
import h5py
import logging
from astropy.table import Table, Column
from astropy import units as u

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("-p","--params", type=str,
                    help = "Parameter file")
parser.add_argument("-q", "--quiet", help = "Suppress extra outputs",
                    action = "store_true")
args = parser.parse_args()
quiet = args.quiet

params_root = re.split(".py", args.params)[0]
if os.path.isfile(params_root+".pyc"):
    os.remove(params_root+".pyc")

import importlib
try:
    params = importlib.import_module(params_root)
    print('Successfully loaded "{0}" as params'.format(args.params))
    importlib.reload(params)
except:
    print('Failed to load "{0}" as params'.format(args.params))
    raise

if quiet:
    quietprint = lambda *a: None
else:
    def quietprint(*args):
        for arg in args:
            print(arg, end=' ')
        print()

# Fitting function definition for later use by Processess

def galaxyFit(inputQueue, printQueue, printlock):
    for gal in iter(inputQueue.get, 'STOP'):
        j = np.argmin(np.abs(z-zobs[gal])) # Find closest model redshift

        flux_obs = obs[gal,:]
        flux_err = obs_err[gal,:]

        #flux_obs[fo <= 0.] = 0.       # Set negative fluxes to zero
        I = np.where(flux_err > 0.)[0] # Find bands with no observation

        if len(I) == 0:
            if include_rest:
                M_scaled = np.ones(len(fo)) * -99.
                restframe_output = ' '.join(M_scaled.astype('str'))
                output_string = '{0} {1} {2} {3} {4} {5} {6} {7}' \
                                ' {8} {9} {10} {11} {12} {13} {14} {15} {16}'.format(gal+1,ID[gal],zobs[gal],-99,-99,-99,-99,-99,-99, -99, -99,-99,len(I),-99,z[j],restframe_output,'\n')
            else:
                output_string = '{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10} {11} {12} {13} {14}'.format(gal+1,ID[gal],zobs[gal],-99,-99,-99,-99,-99,-99,-99, -99,-99,len(I),-99,'\n')
            printQueue.put(output_string)
            continue

        flux_obs = flux_obs[I]                    # and exclude from fit
        flux_err = flux_err[I]
        flux_models = f[j,I,:]

        tot_err = np.sqrt(flux_err**2 + (0.1*flux_obs)**2)

        top = 0.
        bottom = 0.

        for i in range(len(flux_obs)):
            top += (flux_models[i,:]*flux_obs[i])/(tot_err[i]**2)
            bottom += (flux_models[i,:]**2)/(tot_err[i]**2)

        scale = top/bottom
        scale = np.reshape(scale, (n_metal, n_tg, n_tau, n_tauv, n_fesc))

        chisq = 0.
        for i in range(len(flux_obs)):
            chisq += ((np.abs(scale*flux_models[i,:]-flux_obs[i])**2)/(flux_err[i])**2)

        chimin, minind = np.nanmin(chisq), np.nanargmin(chisq)

        if np.isinf(chimin) or np.isnan(minind):
            if include_rest:
                M_scaled = np.ones(len(flux_obs)) * -99.
                restframe_output = ' '.join(M_scaled.astype('str'))
                output_string = '{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10} {11} {12} {13} {14} {15} {16}'.format(gal+1,ID[gal],zobs[gal],-99,-99,-99,-99,-99,-99, -99, -99,-99,len(I),-99,z[j],restframe_output,'\n')
            else:
                output_string = '{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10} {11} {12} {13} {14}'.format(gal+1,ID[gal],zobs[gal],-99,-99,-99,-99,-99,-99,-99, -99,-99,len(I),-99,'\n')
            printQueue.put(output_string)
            continue

        #Find the coordinate of the model with the bestfit mass
        mi, tgi, ti, tvi, fi = np.unravel_index(minind,
                                                   (n_metal, n_tg,
                                                   n_tau, n_tauv, n_fesc))

        Bestfit_Mass = np.log10(scale[mi, tgi, ti, tvi, fi]*flux_corr)
        Bestfit_SFR = (scale[mi, tgi, ti, tvi, fi] *
                       SFR[mi, tgi, ti, tvi, fi]*flux_corr)
        #Bestfit_Beta = beta[tgi,tvi,ti,mi]
        Bestfit_Beta = -99.

        #Scale the observed tot_mag band of the template to be the same as the observed tot_mag band of the galaxy
        #Convert the templates so they are no longer units of per stellar mass

        F_rest = f[0,:]*scale[mi, tgi, ti, tvi, fi]*flux_corr
        restframeMags = 23.9 - 2.5*np.log10(F_rest)

        #UV_rest = UV_flux[0]*scale[tgi,tvi,ti,mi]*flux_corr
        #restframeMUV = 23.9 - 2.5*np.log10(UV_rest)

        M_scaled = restframeMags[:, mi, tgi, ti, tvi, fi]
        #MUV_scaled = restframeMUV[tgi,tvi,ti,mi]
        MUV_scaled = -99.

        if np.isnan(Bestfit_Mass) or np.isinf(chimin):
            Bestfit_Mass = -99
            #M_scaled[:] = -99
            tgs = -99
            tvs = -99
            taus = -99
            mis = -99
            escape_fraction = -99

        else:
            tgs = tg[tgi]/1e9
            tvs = tv[tvi]
            taus = tau[ti]
            mis = metallicities[mi]
            escape_fraction = fesc[fi]

        printlock.acquire()

        print('{0:6d} {1:8d} {2:>5.2f} {3:>7.2f} {4:>8.1f} {5:>8.3f} {6:>5.1f} {7:>8.2f} {8:>4.2f} {9:>5.2f}'.format(gal+1,ID[gal], zobs[gal],Bestfit_Mass,chimin,tgs,tvs,taus,mis,np.log10(Bestfit_SFR)))

        if include_rest:
            restframe_output = ' '.join(M_scaled.astype('str'))
            output_string = '{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10} {11} {12} {13} {14} {15} {16}'.format(gal+1,ID[gal],zobs[gal],Bestfit_Mass,chimin,tgs,tvs,taus,mis, MUV_scaled, minind,Bestfit_SFR,len(I),Bestfit_Beta,z[j],restframe_output,'\n')
        else:
            output_string = '{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10} {11} {12} {13} {14}'.format(gal+1,ID[gal],zobs[gal],Bestfit_Mass,chimin,tgs,tvs,taus,mis, MUV_scaled, minind,Bestfit_SFR,len(I),Bestfit_Beta,'\n')

        printlock.release()
        printQueue.put(output_string)

def galaxyFit2(inputQueue, printQueue, printlock):
    for gal in iter(inputQueue.get, 'STOP'):

        output_string = '{0[0]} {0[1]} {0[2]} {0[3]} {0[4]} {0[5]} ' + \
                        '{0[6]} {0[7]} {0[8]} {0[9]} {0[10]} {0[11]} ' + \
                        '{0[12]} {0[13]} {0[14]}'

        j = np.argmin(np.abs(z-zobs[gal])) # Find closest model redshift

        log_mass_min, log_mass_max = 7, 13
        log_sfr_min, log_sfr_max = -3, 4

        flux_obs = obs[gal,:]
        flux_err = obs_err[gal,:]

        #flux_obs[fo <= 0.] = 0.       # Set negative fluxes to zero
        I = np.where(flux_err > 0.)[0] # Find bands with no observation

        if len(I) == 0:
            output_array = [gal+1, ID[gal], zobs[gal], z[j],
                            -99, -99, -99, -99, -99, -99, -99,
                            -99,-99,len(I),-99,'\n']
            output = output_string.format(output_array)

            if include_rest:
                M_scaled = np.ones(len(flux_obs)) * -99.
                restframe_output = ' '.join(M_scaled.astype('str'))
                output = output + restframe_output + ' \n'

            else:
                output = output + ' \n'
            printQueue.put(output_string)
            continue

        flux_obs = flux_obs[I]                    # and exclude from fit
        flux_err = flux_err[I]
        flux_models = f[j,I,:]

        tot_err = np.sqrt(flux_err**2 + (params.flux_err*flux_obs)**2)

        top = 0.
        bottom = 0.

        for i in range(len(flux_obs)):
            top += (flux_models[i,:]*flux_obs[i])/(tot_err[i]**2)
            bottom += (flux_models[i,:]**2)/(tot_err[i]**2)

        scale = top/bottom
        scale = np.reshape(scale, (n_metal, n_tg, n_tau, n_tauv, n_fesc))

        chisq = 0.
        for i in range(len(flux_obs)):
            chisq += ((np.abs(scale*flux_models[i,:]-flux_obs[i])**2)/(tot_err[i])**2)

        chimin, minind = np.nanmin(chisq), np.nanargmin(chisq)
        likelihood = np.reshape(np.exp(-0.5*chisq),
                                (n_metal, n_tg, n_tau, n_tauv, n_fesc))
        likelihood[np.isnan(likelihood)] = 0.
        likelihood = np.abs(likelihood/likelihood.sum())

        if np.isinf(chimin) or np.isnan(minind):
            output_array = [gal+1, ID[gal], zobs[gal], z[j],
                            -99, -99, -99, -99, -99, -99, -99,
                            -99,-99,len(I),-99,'\n']
            output = output_string.format(output_array)

        else:
            #Find the coordinate of the model with the bestfit mass
            mi, tgi, ti, tvi, fi = np.unravel_index(minind,
                                                       (n_metal, n_tg,
                                                       n_tau, n_tauv, n_fesc))

            Masses = np.abs(np.log10(scale*flux_corr))
            SFRs = np.abs(np.log10(scale * SFR * flux_corr))

            mass_hist = np.histogram(Masses.flatten(),
                                     range = (log_mass_min, log_mass_max),
                                     bins = 120,
                                     weights = likelihood.flatten(),
                                     density = True)

            sfr_hist = np.histogram(SFRs.flatten(),
                                     range = (log_sfr_min, log_sfr_max),
                                     bins = 140,
                                     weights = likelihood.flatten(),
                                     density = True)

            Bestfit_Mass = np.abs(np.log10(scale[mi, tgi, ti, tvi, fi]*flux_corr))
            Bestfit_SFR = np.abs(np.log10(scale[mi, tgi, ti, tvi, fi] *
                                   SFR[mi, tgi, ti, tvi, fi]*flux_corr))

            if np.isnan(Bestfit_Mass) or np.isinf(chimin):
                Bestfit_Mass = -99
                #M_scaled[:] = -99
                tgs = -99
                tvs = -99
                taus = -99
                mis = -99
                escape_fraction = -99

            else:
                tgs = tg[tgi]/1e9
                tvs = tv[tvi]
                taus = tau[ti]
                mis = metallicities[mi]
                escape_fraction = fesc[fi]

            m16, m50, m84 = weighted_quantile(Masses.flatten(),
                                              [0.16, 0.5, 0.84],
                                              sample_weight=likelihood.flatten(),
                                              values_sorted=False)
            s16, s50, s84 = weighted_quantile(SFRs.flatten(),
                                              [0.16, 0.5, 0.84],
                                              sample_weight=likelihood.flatten(),
                                              values_sorted=False)

            printlock.acquire()

            MUV_scaled = -99.
            Bestfit_Beta = -99.

            print_string = "{0[0]:6d} {0[1]:8d} {0[2]:>5.2f} " + \
                           "{0[3]:>7.2f} {0[4]:>8.1f} {0[5]:>8.3f} " + \
                           "{0[6]:>5.1f} {0[7]:>8.2f} {0[8]:>4.2f} " + \
                           "{0[9]:>5.2f}"

            print_array = [gal+1, ID[gal], zobs[gal],
                           Bestfit_Mass, chimin,
                           tgs, tvs, taus, mis,
                           Bestfit_SFR]
            print(print_string.format(print_array))

            output_string = '{n} {id} {zobs} {ztemp} {mass_best} {sfr_best} '+ \
                            '{chi_best} {tg} {tvs} {taus} {mis} {fesc} '+ \
                            '{mass_med} {mass_l68} {mass_u68} ' + \
                            '{sfr_med} {sfr_l68} {sfr_u68} ' + \
                            '{nfilts} '

            output_values = {'n': gal+1,
                             'id': ID[gal],
                             'zobs': zobs[gal], 'ztemp':z[j],
                             'mass_best': Bestfit_Mass,
                             'sfr_best': Bestfit_SFR,
                             'chi_best': chimin,
                             'tg': tgs, 'tvs': tvs, 'taus': taus,
                             'mis': mis, 'fesc': escape_fraction,
                             'mass_med': m50, 'mass_l68': m16, 'mass_u68': m84,
                             'sfr_med': s50, 'sfr_l68': s16, 'sfr_u68': s84,
                             'nfilts': len(I)}

            output_array = [gal+1, ID[gal], zobs[gal],
                            Bestfit_Mass, chimin, tgs, tvs, taus, mis,
                            MUV_scaled, minind, Bestfit_SFR, len(I), -99., '\n']
            output = output_string.format(**output_values)

        if include_rest:
            if np.isinf(chimin) or np.isnan(minind):
                M_scaled = np.ones(len(flux_obs)) * -99.
                restframe_output = ' '.join(M_scaled.astype('str'))
                output = output + restframe_output + ' \n'

            else:
                F_rest = np.array(f[0, :, mi, tgi, ti, tvi, fi] *
                                  scale[mi, tgi, ti, tvi, fi] * flux_corr)
                restframeMags = 23.9 - 2.5*np.log10(F_rest)
                restframe_output = ' '.join(restframeMags.astype('str'))
                output = output + restframe_output + ' \n'
        else:
            output = output + ' \n'

        printlock.release()
        printQueue.put([output, mass_hist, sfr_hist])

def galaxyFitPlus(inputQueue, printQueue, printlock):
    for gal in iter(inputQueue.get, 'STOP'):
        mass_range = 7, 13
        log_sfr_min, log_sfr_max = -3, 4

        j = np.argmin(np.abs(z-zobs[gal])) # Find closest model redshift

        fo = obs[gal,:]
        ferr = obs_err[gal,:]

        flux_obs[fo <= 0.] = 0.       # Set negative fluxes to zero
        #print fo
        I = (ferr > 0.)*(ferr < 1e6) # Find bands with no observation
        fo = flux_obs[I]                    # and exclude from fit
        ferr = flux_err[I]
        fm = f[I,j,:]
        #print flux_models[:,0,0,0,0]

        top = 0.
        bottom = 0.

        for i in range(len(fo)):
            top += (flux_models[i,:]*flux_obs[i])/(flux_err[i]**2)
            bottom += (flux_models[i,:]**2)/(flux_err[i]**2)

        scale = top/bottom
        scale = np.reshape(scale, (n_metal, n_tg, n_tau, n_tauv, n_fesc))

        chisq = 0.
        for i in range(len(fo)):
            chisq += ((np.abs(scale*flux_models[i,:]-flux_obs[i])**2)/(flux_err[i])**2)

        chimin, minind = np.nanmin(chisq), np.nanargmin(chisq)

        chisq -= (chisq.min() - 1)
        likelihood = np.exp(-0.5*chisq)
        likelihood /= likelihood.sum()

        if np.isinf(chimin) or np.isnan(minind) or len(fo) == 0:
            output_string = '{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} \
            {10} {11} {12} {13} {14} {15} {16} {17} {18}'.format(gal+1,ID[gal],zobs[gal],
                                                                 -99,-99,-99,-99,-99,-99,
                                                                 -99, -99, -99, -99,-99,-99,-99,
                                                                 len(I),-99,'\n')

            massLikelihood = np.zeros(mass_bins+1)
            massLikelihood[0] = gal
            muvLikelihood = np.zeros(muv_bins+1)
            muvLikelihood[0] = gal
            betaLikelihood = np.zeros(beta_bins+1)
            betaLikelihood[0] = gal
            #tauLikelihood = np.zeros(n_tau)
            #tauLikelihood = np.insert(tauLikelihood,0,gal)
            printQueue.put([output_string,massLikelihood,muvLikelihood,betaLikelihood])
            continue

        #Find the coordinate of the model with the bestfit mass
        si,tgi,tvi,ti,mi = np.unravel_index(minind,(mass_bins,n_tg,n_tauv,n_tau,n_ssp))
        Bestfit_Mass = np.log10(mass_range[si]*flux_corr)
        Bestfit_SFR = (mass_range[si]*SFR[tgi,ti,mi]*flux_corr)
        Bestfit_Beta = beta[tgi,tvi,ti,mi]

        F_rest = f[:,0]*mass_range[likelihood.argmax(0)]*flux_corr
        restframeMags = 23.9 - 2.5*np.log10(F_rest)

        UV_rest = UV_flux[0]*mass_range[likelihood.argmax(0)]*flux_corr
        restframeMUV = 23.9 - 2.5*np.log10(UV_rest)

        Bestfit_restframeMags = restframeMags[:,tgi,tvi,ti,mi]
        Bestfit_restframeMUV = restframeMUV[tgi,tvi,ti,mi]

        if np.isnan(Bestfit_Mass) or np.isinf(chimin):
            Bestfit_Mass = -99
            #M_scaled[:] = -99
            tgs = -99
            tvs = -99
            taus = -99
            mis = -99

        else:
            tgs = tg[tgi]/1.e9
            tvs = tv[tvi]
            taus = tau[ti]/1.e9
            mis = mi

        """
        Likelihood array section:
        """
        mass_hist = np.histogram(np.log10(mass_))

        printlock.acquire()

        if calc_mode:
            print('{0:4d} {1:6d} {2:>6.2f} {3:>8.1f} {4:>6.2f}'.format(gal+1,ID[gal],Bestfit_Mass,chimin, np.log10(Mode_Mass), '/n'))
        else:
            print('{0:6d} {1:8f} {2:>5.2f} {3:>7.2f} {4:>8.1f} {5:>8.3f} {6:>5.1f} {7:>8.2f} {8:>3d} {9:>5.2f}'.format(gal+1,int(ID[gal]),zobs[gal],Bestfit_Mass,chimin,tgs,tvs,taus,mis,np.log10(Bestfit_SFR)))

        output_string = '{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10} {11} {12} {13} {14} {15}'.format(gal+1,int(ID[gal]),zobs[gal],Bestfit_Mass,chimin,tgs,tvs,taus,mis,Bestfit_restframeMags[tot],Bestfit_restframeMUV,minind,Bestfit_SFR,len(I),Bestfit_Beta,'\n')

        printlock.release()
        printQueue.put([output_string, massLikelihoods, muvLikelihoods, betaLikelihoods])

def getObservations(inputpath):
    input_data = Table.read(inputpath,format=input_format)

    column_names = list(input_data.columns.keys())

    ID = input_data[ID_col]
    zobs = input_data[z_col]

    filter_names = []

    k,l = 0,0
    for ii in range(len(column_names)):
        if column_names[ii].lower().endswith(flux_col_end.lower()):
            if k == 0:
                fluxes = input_data[column_names[ii]]
            else:
                fluxes = np.column_stack((fluxes,input_data[column_names[ii]]))
            k+=1
            filter_names.append(column_names[ii])

        if column_names[ii].lower().endswith(fluxerr_col_end.lower()):
            if l == 0:
                fluxerrs = input_data[column_names[ii]]
            else:
                fluxerrs = np.column_stack((fluxerrs,input_data[column_names[ii]]))
            l+=1
    """
    if filts_used != None:
        try:
            fluxes = fluxes[:,filts_used]
            fluxerrs = fluxerrs[:,filts_used]
        except:r
            print('Filter mismatch 1')
            # Array slicing fail
    """
    return ID, zobs, fluxes, fluxerrs, k, filter_names

class _function_wrapper(object):
    """
    This is a hack to make the likelihood function pickleable when ``args``
    or ``kwargs`` are also included.

    Stolen from emcee
    """
    def __init__(self, f, args, kwargs):
        self.f = f
        self.args = args
        self.kwargs = kwargs

    def __call__(self, x):
        try:
            return self.f(x, *self.args, **self.kwargs)
        except:
            import traceback
            print("emcee: Exception while calling your likelihood function:")
            print("  params:", x)
            print("  args:", self.args)
            print("  kwargs:", self.kwargs)
            print("  exception:")
            traceback.print_exc()
            raise

def weighted_quantile(values, quantiles, sample_weight=None, values_sorted=False, old_style=False):
    """ Very close to np.percentile, but supports weights.
    NOTE: quantiles should be in [0, 1]!
    :param values: np.array with data
    :param quantiles: array-like with many quantiles needed
    :param sample_weight: array-like of the same length as `array`
    :param values_sorted: bool, if True, then will avoid sorting of initial array
    :param old_style: if True, will correct output to be consistent with np.percentile.
    :return: np.array with computed quantiles.
    """
    values = np.array(values)
    quantiles = np.array(quantiles)
    if sample_weight is None:
        sample_weight = np.ones(len(values))
    sample_weight = np.array(sample_weight)
    assert np.all(quantiles >= 0) and np.all(quantiles <= 1), 'quantiles should be in [0, 1]'

    if not values_sorted:
        sorter = np.argsort(values)
        values = values[sorter]
        sample_weight = sample_weight[sorter]

    weighted_quantiles = np.cumsum(sample_weight) - 0.5 * sample_weight
    if old_style:
        # To be convenient with np.percentile
        weighted_quantiles -= weighted_quantiles[0]
        weighted_quantiles /= weighted_quantiles[-1]
    else:
        weighted_quantiles /= np.sum(sample_weight)
    return np.interp(quantiles, weighted_quantiles, values)

if __name__ == '__main__':

    logfile = open("error.log", "w")
    original_stderr = sys.stderr
    sys.stderr = logfile

    start = time.time()

    """
    SECTION 1

    """
    model_path = params.model_path

    input_catalog = params.input_catalog
    input_format = params.input_format
    z_col = params.z_col
    ID_col = params.ID_col
    flux_col_end = params.flux_col_end
    fluxerr_col_end = params.fluxerr_col_end

    ncpus = params.ncpus
    filts_used = params.filts_used
    include_rest = params.include_rest

    output_path = params.output_catalog_path
    output_format = params.output_format
    output_hdf_path = params.output_hdf_path

    calc_mode = params.fitting_mode
    flux_corr = params.flux_corr

    ID, zobs, obs, obs_err, filters_found, filter_names = getObservations(input_catalog)

    """
    Section 2

    """

    print("Loading synthetic mags and mass array:")
    models = h5py.File(model_path, 'r')
    tg = models['ages'].value
    tv = models['dust'].value
    tau = models['sfh'].value
    metallicities = models['metallicities'].value
    fesc = models['fesc'].value

    Mshape = models['fluxes'].shape
    z = models['z']
    nfilts = Mshape[1]
    n_metal = Mshape[2]
    n_tg = Mshape[3]
    n_tau = Mshape[4]
    n_tauv = Mshape[5]
    n_fesc = Mshape[6]

    #UV_flux = synmags['UV_flux']
    SFR = models['SFR']
    Ms = models['Ms']

    if (nfilts == filters_found) and (filts_used == None):
        f = models['fluxes']

    elif filts_used != None:
        try:
            f = models['fluxes'][:,filts_used]
            obs = obs[:,filts_used]
            obs_err = obs_err[:,filts_used]
            filter_names = np.array(filter_names)[filts_used]
        except:
            print('Mis-match between model and observed filter numbers')
            raise
            # Slice fail

    print ("Done.")

    """
    SECTION 3
    """
    if os.path.isfile(output_path+".temp_output.txt"):
        os.remove(output_path+".temp_output.txt")
    temp_file = open(output_path+".temp_output.txt","w")

    """
    SECTION 4
    Chi-sq calculation

    """
    out_string = '{0:6s} {1:8s} {2:>5s} {3:>7s} {4:>8s} {5:>8s}' + \
                 '{6:>5s} {7:>8s} {8:>4s} {9:>5s}'

    print(out_string.format('N','ID','zobs','Best', 'chimin',
                            'tg', 'tauv','tau','met', 'sfr'))

    loop_start = time.time()
    ncpus = np.clip(ncpus, 1, multiprocessing.cpu_count())

    inputQueue = multiprocessing.Queue()
    printQueue = multiprocessing.Queue()
    printlock = multiprocessing.Lock()

    if calc_mode == 'hist':
        output_hdf = h5py.File(output_hdf_path, 'w')
        output_hdf.create_dataset("mass_pdf", (len(ID), 120), dtype="f")
        output_hdf.create_dataset("sfr_pdf", (len(ID), 140), dtype="f")
        fitFunction = galaxyFit2

    else:
        fitFunction = galaxyFit

    for i in range( ncpus ):
        multiprocessing.Process(target = fitFunction,
                                args = (inputQueue, printQueue,
                                        printlock)).start()

    # Put elements in the send queue for processing
    for gal in range( len(ID) ):
        inputQueue.put( gal )

    if calc_mode == 'hist':
        for i, gal in enumerate(ID):
            printout, mass_hist, sfr_hist = printQueue.get()
            if i == 0:
                mass_centers = 0.5*(mass_hist[1][1:] + mass_hist[1][:-1])
                sfr_centers = 0.5*(sfr_hist[1][1:] + sfr_hist[1][:-1])

                output_hdf.create_dataset("mass_bins", data = mass_centers)
                output_hdf.create_dataset("sfr_bins", data = sfr_centers)
            output_hdf["mass_pdf"][i] = mass_hist[0]
            output_hdf["sfr_pdf"][i] = sfr_hist[0]
            temp_file.write( printout )
        #tau_array.tofile(tau_file)
    else:
        for i, gal in enumerate(ID):
            printout = printQueue.get()
            temp_file.write( printout )
            #print len(mass_array), len(muv_array), len(beta_array)

    # Stop all the running processes
    for i in range( ncpus ):
        inputQueue.put( 'STOP' )

    # Close both send and receive queues
    inputQueue.close()
    printQueue.close()

    temp_file.close()
    models.close()
    output_hdf.close()
    print("Fitting time taken: {0:.2f} {1}".format(time.time()-loop_start,
                                                   '\n'))

    """
    Section 3
    Reload, format and save output table
    """
    while temp_file.closed == False:
        pause(0.1)

    data = np.loadtxt(output_path+".temp_output.txt")
    try:
        rows, cols = data.shape
    except:
        cols = len(data)

    output = Table()

    names = ['N', 'ID', 'z', 'zmodel',
             'Mass_best', 'SFR_best', 'chi_best',
             'Age_best','Dust_best', 'SFH_best',
             'Metallicity_best', 'fesc_best',
             'Mass_median', 'Mass_l68', 'Mass_u68',
             'SFR_median', 'SFR_l68', 'SFR_u68',
             'Nfilts']

    units = [None, None, None, None,
             u.Msun, u.Msun/u.yr, None,
             u.Gyr, None, None,
             None, None,
             u.Msun, u.Msun, u.Msun,
             u.Msun/u.yr, u.Msun/u.yr, u.Msun/u.yr,
             None]

    types = ['i4', 'i4', 'f4', 'f4',
             'f4', 'f4', 'f4',
             'f4', 'f4', 'f4',
             'f4', 'f4',
             'f4', 'f4', 'f4',
             'f4', 'f4', 'f4',
             'i4']

    if include_rest:
        for name in filter_names:
            names.append(name[:-len(flux_col_end)]+'_rest')
            units.append(u.mag)
            types.append('f4')

    for col in range(cols):
        column = Column( data[:,col], name = names[col], unit=units[col], dtype=types[col])
        output.add_column(column)

    table_format = 'ascii.commented_header'
    output.sort('ID')
    if os.path.isfile(output_path):
        os.remove(output_path)
    output.write(output_path,format=table_format)
    print('Catalog saved')

    os.remove(temp_file.name)

    print()
    print("Total time taken: "+str(time.time()-start))

    sys.stderr = original_stderr
    logfile.close()


#!/usr/bin/env python
# coding: utf-8

# In[1]:

# Copyright 2018 The TensorFlow Authors All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Simple example of contextual bandits simulation.

Code corresponding to:
Deep Bayesian Bandits Showdown: An Empirical Comparison of Bayesian Deep Networks
for Thompson Sampling, by Carlos Riquelme, George Tucker, and Jasper Snoek.
https://arxiv.org/abs/1802.09127
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import time,sys
import tensorflow as tf
from absl import app
from absl import flags
import numpy as np
import os
import matplotlib.pyplot as plt
import matplotlib
# print(matplotlib.get_backend())
from bandits.helpers.benchmarker import Benchmarker
from bandits.algorithms.lin_ucb import LinUCB
from bandits.algorithms.lin_epsilon import LinEpsilon
from bandits.algorithms.neural_lin_ucb import NeuralLinUCB

from bandits.algorithms.bootstrapped_bnn_sampling import BootstrappedBNNSampling
from bandits.core.contextual_bandit import run_contextual_bandit
# from bandits.data.data_sampler import sample_mushroom_data
from bandits.algorithms.fixed_policy_sampling import FixedPolicySampling
from bandits.algorithms.linear_full_posterior_sampling import LinearFullPosteriorSampling
from bandits.algorithms.neural_linear_sampling import NeuralLinearPosteriorSampling
from bandits.algorithms.parameter_noise_sampling import ParameterNoiseSampling
from bandits.algorithms.posterior_bnn_sampling import PosteriorBNNSampling
from bandits.data.synthetic_data_sampler import sample_linear_data
from bandits.data.synthetic_data_sampler import sample_wheel_bandit_data
from bandits.data.environments import  *
from bandits.data.wasserstein_gans import WGANCovertype

from bandits.data.bootstrap_thompson_sampling import generate_uniform_artificial, gan_artificial_covertype,\
    gan_artificial_mushroom, gan_artificial_linear, gan_artificial_wheel

from bandits.algorithms.neural_linear_sampling import NeuralLinearPosteriorSampling
from bandits.data.synthetic_data_sampler import sample_sparse_linear_data
from bandits.data.synthetic_data_sampler import sample_wheel_bandit_data
from bandits.algorithms.uniform_sampling import UniformSampling

base_route = os.getcwd()
data_route = 'contextual_bandits/datasets'

FLAGS = flags.FLAGS
FLAGS.set_default('alsologtostderr', True)

flags.DEFINE_string('logdir', base_route, 'Base directory to save output')
FLAGS(sys.argv)

############# STARTS HERE ##############"""
name = "linear"
test_name = "full_analysis_" + name

if name=="linear":
    num_actions = 8
    context_dim = 10
    num_contexts = 1500
    # noise_stds = [0.01 * (i + 1) for i in range(num_actions)]
    noise_stds = [1 for i in range(num_actions)]

    wgan = WGANCovertype(context_dim, file="linear")
    wgan.train(epochs=400, batch_size=32, sample_interval=50)
    artificial_data_generator = lambda: gan_artificial_linear(wgan, n_samples=50, n_actions=num_actions)

elif name=="mushroom":
    num_actions = 2
    context_dim = 117
    num_contexts = 1500
    from bandits.data.wasserstein_gans import WGANMushroom
    wgan = WGANMushroom(context_dim)
    wgan.train(epochs=2000, batch_size=32, sample_interval=50)
    artificial_data_generator = lambda: gan_artificial_mushroom(wgan, n_samples=50, n_actions=num_actions)

elif name=="wheel":
    num_actions = 5
    context_dim = 2
    num_contexts = 1500
    delta = 0.95
    mean_v = [1.0, 1.0, 1.0, 1.0, 1.2]
    std_v = [0.05, 0.05, 0.05, 0.05, 0.05]
    mu_large = 50
    std_large = 0.01
    from bandits.data.wasserstein_gans import WGANCovertype
    wgan = WGANCovertype(context_dim, file="wheel")
    wgan.train(epochs=1000, batch_size=32, sample_interval=50)
    artificial_data_generator = lambda: gan_artificial_wheel(wgan, n_samples=50, n_actions=num_actions)

elif name=="covertype":
    num_actions = 7
    context_dim = 54
    num_contexts = 1500
    from bandits.data.wasserstein_gans import WGANCovertype
    wgan = WGANCovertype(context_dim)
    wgan.train(epochs=4000, batch_size=32, sample_interval=50)
    artificial_data_generator = lambda: gan_artificial_covertype(wgan, n_samples=50, n_actions=num_actions)
else:
    raise Exception('name not recognized')

def dataset_proto(name=name):
    if name=="linear":
        dataset, _, opt_linear = sample_linear_data(num_contexts, context_dim,
                                                    num_actions, sigma=noise_stds)
    elif name=="wheel":
        dataset, opt_linear = sample_wheel_bandit_data(num_contexts, delta,
                                                      mean_v, std_v,
                                                      mu_large, std_large)

    elif name=="mushroom":
        mush = Mushrooms(num_contexts=num_contexts)
        dataset = mush.table
        opt_rewards, opt_actions = mush.opts[:,0], mush.opts[:,1]
        opt_linear = (opt_rewards, opt_actions)

    elif name=="covertype":
        cov = Covertype(num_contexts=num_contexts)
        dataset = cov.table
        opt_rewards, opt_actions = cov.opts[:,0], cov.opts[:,1]
        opt_linear = (opt_rewards, opt_actions)

    return dataset, opt_linear

print(dataset_proto()[0].shape)

# Params for algo templates
hparams = tf.contrib.training.HParams(num_actions=num_actions)

hparams_linear = tf.contrib.training.HParams(num_actions=num_actions,
                                             context_dim=context_dim,
                                             a0=6,
                                             b0=6,
                                             lambda_prior=0.25,
                                             initial_pulls=2)

hparams_rms = tf.contrib.training.HParams(num_actions=num_actions,
                                        context_dim=context_dim,
                                        init_scale=0.3,
                                        activation=tf.nn.relu,
                                        layer_sizes=[50],
                                        batch_size=512,
                                        activate_decay=True,
                                        initial_lr=0.1,
                                        max_grad_norm=5.0,
                                        show_training=False,
                                        freq_summary=1000,
                                        buffer_s=-1,
                                        initial_pulls=2,
                                        optimizer='RMS',
                                        reset_lr=True,
                                        lr_decay_rate=0.5,
                                        training_freq=50,
                                        training_epochs=50,
                                        bootstrap=None)

hparams_rms_bootstrapped = tf.contrib.training.HParams(num_actions=num_actions,
                                        context_dim=context_dim,
                                        init_scale=0.3,
                                        activation=tf.nn.relu,
                                        layer_sizes=[50],
                                        batch_size=512,
                                        activate_decay=True,
                                        initial_lr=0.1,
                                        max_grad_norm=5.0,
                                        show_training=False,
                                        freq_summary=1000,
                                        buffer_s=-1,
                                        initial_pulls=2,
                                        optimizer='RMS',
                                        reset_lr=True,
                                        lr_decay_rate=0.5,
                                        training_freq=50,
                                        training_epochs=50,
                                        bootstrap=artificial_data_generator)

hparams_rmsb = tf.contrib.training.HParams(num_actions=num_actions,
                                        context_dim=context_dim,
                                        init_scale=0.3,
                                        activation=tf.nn.relu,
                                        layer_sizes=[50],
                                        batch_size=512,
                                        activate_decay=True,
                                        initial_lr=0.1,
                                        max_grad_norm=5.0,
                                        show_training=False,
                                        freq_summary=1000,
                                        buffer_s=-1,
                                        initial_pulls=2,
                                        optimizer='RMS',
                                        reset_lr=True,
                                        lr_decay_rate=0.5,
                                        training_freq=50,
                                        training_epochs=50,
                                        bootstrap=None,
                                        q=3,p=0.95)

hparams_rmsb_bootstrapped = tf.contrib.training.HParams(num_actions=num_actions,
                                        context_dim=context_dim,
                                        init_scale=0.3,
                                        activation=tf.nn.relu,
                                        layer_sizes=[50],
                                        batch_size=512,
                                        activate_decay=True,
                                        initial_lr=0.1,
                                        max_grad_norm=5.0,
                                        show_training=False,
                                        freq_summary=1000,
                                        buffer_s=-1,
                                        initial_pulls=2,
                                        optimizer='RMS',
                                        reset_lr=True,
                                        lr_decay_rate=0.5,
                                        training_freq=50,
                                        training_epochs=50,
                                        q=3,p=0.95,
                                        bootstrap=artificial_data_generator)

hparams_dropout = tf.contrib.training.HParams(num_actions=num_actions,
                                        context_dim=context_dim,
                                        init_scale=0.3,
                                        use_dropout=True,
                                        keep_prob=0.95,
                                        activation=tf.nn.relu,
                                        layer_sizes=[50],
                                        batch_size=512,
                                        activate_decay=True,
                                        initial_lr=0.1,
                                        max_grad_norm=5.0,
                                        show_training=False,
                                        freq_summary=1000,
                                        buffer_s=-1,
                                        initial_pulls=2,
                                        optimizer='RMS',
                                        reset_lr=True,
                                        lr_decay_rate=0.5,
                                        training_freq=50,
                                        training_epochs=50,
                                        bootstrap=None)

hparams_dropout_bootstrapped = tf.contrib.training.HParams(num_actions=num_actions,
                                        context_dim=context_dim,
                                        init_scale=0.3,
                                        use_dropout=True,
                                        keep_prob=0.95,
                                        activation=tf.nn.relu,
                                        layer_sizes=[50],
                                        batch_size=512,
                                        activate_decay=True,
                                        initial_lr=0.1,
                                        max_grad_norm=5.0,
                                        show_training=False,
                                        freq_summary=1000,
                                        buffer_s=-1,
                                        initial_pulls=2,
                                        optimizer='RMS',
                                        reset_lr=True,
                                        lr_decay_rate=0.5,
                                        training_freq=50,
                                        training_epochs=50,
                                        bootstrap=artificial_data_generator)

hparams_linucb = tf.contrib.training.HParams(num_actions=num_actions,
                                        context_dim=context_dim,
                                        alpha=1,
                                        lam=0.1)

hparams_neural_linucb = tf.contrib.training.HParams(num_actions=num_actions,
                                        context_dim=context_dim,
                                        init_scale=0.3,
                                        activation=tf.nn.relu,
                                        layer_sizes=[50],
                                        batch_size=512,
                                        activate_decay=True,
                                        initial_lr=0.1,
                                        max_grad_norm=5.0,
                                        show_training=False,
                                        freq_summary=1000,
                                        buffer_s=-1,
                                        initial_pulls=2,
                                        optimizer='RMS',
                                        reset_lr=True,
                                        lr_decay_rate=0.5,
                                        training_freq=50,
                                        training_epochs=50,
                                        training_freq_network=50,
                                        bootstrap=None,
                                        alpha=1,
                                        lam=0.1)
hparams_neural_linucb_bootstrapped = tf.contrib.training.HParams(num_actions=num_actions,
                                        context_dim=context_dim,
                                        init_scale=0.3,
                                        activation=tf.nn.relu,
                                        layer_sizes=[50],
                                        batch_size=512,
                                        activate_decay=True,
                                        initial_lr=0.1,
                                        max_grad_norm=5.0,
                                        show_training=False,
                                        freq_summary=1000,
                                        buffer_s=-1,
                                        initial_pulls=2,
                                        optimizer='RMS',
                                        reset_lr=True,
                                        lr_decay_rate=0.5,
                                        training_freq=50,
                                        training_epochs=50,
                                        training_freq_network=50,
                                        bootstrap=artificial_data_generator,
                                        alpha=1,
                                        lam=0.1)
hparams_neural_linthomson = tf.contrib.training.HParams(num_actions=num_actions,
                                        context_dim=context_dim,
                                        init_scale=0.3,
                                        activation=tf.nn.relu,
                                        layer_sizes=[50],
                                        batch_size=512,
                                        activate_decay=True,
                                        initial_lr=0.1,
                                        max_grad_norm=5.0,
                                        show_training=False,
                                        freq_summary=1000,
                                        buffer_s=-1,
                                        optimizer='RMS',
                                        reset_lr=True,
                                        lr_decay_rate=0.5,
                                        training_freq=50,
                                        training_epochs=50,
                                        training_freq_network=50,
                                        bootstrap=None,
                                        a0=6,
                                        b0=6,
                                        lambda_prior=0.25,
                                        initial_pulls=2)

hparams_neural_linthomson_bootstrapped = tf.contrib.training.HParams(num_actions=num_actions,
                                        context_dim=context_dim,
                                        init_scale=0.3,
                                        activation=tf.nn.relu,
                                        layer_sizes=[50],
                                        batch_size=512,
                                        activate_decay=True,
                                        initial_lr=0.1,
                                        max_grad_norm=5.0,
                                        show_training=False,
                                        freq_summary=1000,
                                        buffer_s=-1,
                                        optimizer='RMS',
                                        reset_lr=True,
                                        lr_decay_rate=0.5,
                                        training_freq=50,
                                        training_epochs=50,
                                        training_freq_network=50,
                                        bootstrap=artificial_data_generator,
                                        a0=6,
                                        b0=6,
                                        lambda_prior=0.25,
                                        initial_pulls=2)

hparams_pnoise = tf.contrib.training.HParams(num_actions=num_actions,
                                       context_dim=context_dim,
                                       init_scale=0.3,
                                       activation=tf.nn.relu,
                                       layer_sizes=[50],
                                       batch_size=512,
                                       activate_decay=True,
                                       initial_lr=0.1,
                                       max_grad_norm=5.0,
                                       show_training=False,
                                       freq_summary=1000,
                                       buffer_s=-1,
                                       initial_pulls=2,
                                       optimizer='RMS',
                                       reset_lr=True,
                                       lr_decay_rate=0.5,
                                       training_freq=50,
                                       training_epochs=100,
                                       noise_std=0.05,
                                       eps=0.1,
                                       d_samples=300
                                      )

hparams_pnoise_bootstrapped = tf.contrib.training.HParams(num_actions=num_actions,
                                       context_dim=context_dim,
                                       init_scale=0.3,
                                       activation=tf.nn.relu,
                                       layer_sizes=[50],
                                       batch_size=512,
                                       activate_decay=True,
                                       initial_lr=0.1,
                                       max_grad_norm=5.0,
                                       show_training=False,
                                       freq_summary=1000,
                                       buffer_s=-1,
                                       initial_pulls=2,
                                       optimizer='RMS',
                                       reset_lr=True,
                                       lr_decay_rate=0.5,
                                       training_freq=50,
                                       training_epochs=100,
                                       noise_std=0.05,
                                       eps=0.1,
                                       d_samples=300,
                                       bootstrap=artificial_data_generator
                                      )

hparams_lineps = tf.contrib.training.HParams(num_actions=num_actions,
                                        context_dim=context_dim,
                                        lam=0.1,
                                        eps=0.05)

random_proto = lambda : UniformSampling('Uniform Sampling', hparams)
neural_greedy_proto = lambda : PosteriorBNNSampling('NeuralGreedy', hparams_rms, 'RMSProp')
neural_greedy_proto_bootstrapped = lambda : PosteriorBNNSampling('NeuralGreedy_artificial_data', hparams_rms_bootstrapped, 'RMSProp')

bootstrap_proto = lambda : BootstrappedBNNSampling('BootRMS', hparams_rmsb)
bootstrap_proto_bootstrapped = lambda : BootstrappedBNNSampling('BootRMS_artificial_data', hparams_rmsb_bootstrapped)

noise_proto = lambda : ParameterNoiseSampling('ParamNoise', hparams_pnoise)
noise_proto_bootstrapped = lambda : ParameterNoiseSampling('ParamNoise_artificial_data', hparams_pnoise_bootstrapped)

dropout_proto = lambda : PosteriorBNNSampling('Dropout', hparams_dropout, 'RMSProp')
dropout_proto_bootstrapped = lambda : PosteriorBNNSampling('Dropout_artificial_data', hparams_dropout_bootstrapped, 'RMSProp')

linThompson_proto = lambda : LinearFullPosteriorSampling('linThompson', hparams_linear)
linUCB_proto = lambda : LinUCB('linUCB', hparams_linucb)
linEps_proto = lambda : LinEpsilon('LinEpsilon', hparams_lineps)

neuralLinUCB_proto = lambda : NeuralLinUCB('NeuralLinUCB', hparams_neural_linucb, 'RMSProp')
neuralLinThomson_proto = lambda : NeuralLinearPosteriorSampling('NeuralLinThomson', hparams_neural_linthomson, 'RMSProp')
neuralLinUCB_proto_bootstrapped = lambda : NeuralLinUCB('NeuralLinUCB_artificial_data', hparams_neural_linucb_bootstrapped, 'RMSProp')
neuralLinThomson_proto_bootstrapped = lambda : NeuralLinearPosteriorSampling('NeuralLinThomson_artificial_data', hparams_neural_linthomson_bootstrapped, 'RMSProp')

algo_protos = [linUCB_proto,
    neuralLinUCB_proto, neuralLinUCB_proto_bootstrapped,
    dropout_proto, dropout_proto_bootstrapped,
    bootstrap_proto, bootstrap_proto_bootstrapped,
    noise_proto, noise_proto_bootstrapped,
    neuralLinThomson_proto, neuralLinThomson_proto_bootstrapped,
    linEps_proto,
    linThompson_proto,
    neural_greedy_proto, neural_greedy_proto_bootstrapped,
    random_proto]

# Run experiments several times save and plot results
benchmarker = Benchmarker(algo_protos, dataset_proto, num_actions, context_dim, nb_contexts=num_contexts, test_name=test_name)

benchmarker.run_experiments(50)
benchmarker.save_results('./results/')
benchmarker.save_final_res_to_tex('./results/')
benchmarker.display_results(save_path='./results/')


# -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2017-2018 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.

"""
Module exports :class:`MunsonThurber1997`
               :class:`MunsonThurber1997Vector`.
"""
import numpy as np

from openquake.hazardlib.gsim.base import GMPE
from openquake.hazardlib import const
from openquake.hazardlib.imt import PGA

class MunsonThurber1997(GMPE):
    """
    Implements GMPE developed by Clifford G. Munson and Clifford H. Thurber
    and published as "Analysis of the Attenuation of Strong Ground Motion
    on the Island of Hawaii" (1997, Bulletin of the Seismological Society
    of America, Vol. 87, No. 4, pp. 954-960).
    """

    #: Supported tectonic region type is volcanic,
    #: see paragraph 'Introduction', page 99.
    DEFINED_FOR_TECTONIC_REGION_TYPE = const.TRT.VOLCANIC

    #: Supported intensity measure types is spectral acceleration,
    #: see table 3, pag. 110
    DEFINED_FOR_INTENSITY_MEASURE_TYPES = set([PGA])

    #: Supported intensity measure component is maximum horizontal
    #: :attr:`~openquake.hazardlib.const.IMC.VECTORIAL`,
    DEFINED_FOR_INTENSITY_MEASURE_COMPONENT = const.IMC.HORIZONTAL

    #: Supported standard deviation type is total
    DEFINED_FOR_STANDARD_DEVIATION_TYPES = set([
        const.StdDev.TOTAL
    ])

    #: Required site parameters is Vs30.
    #: See paragraph 'Predictor Variables', pag 103
    REQUIRES_SITES_PARAMETERS = set(('vs30', ))

    #: Required rupture parameter is magnitude
    REQUIRES_RUPTURE_PARAMETERS = set(('mag', ))

    #: Required distance measure is hypocentral distance
    #: see page 18 in Atkinson and Boore's manuscript
    REQUIRES_DISTANCES = set(('rjb', ))

    def get_mean_and_stddevs(self, sites, rup, dists, imt, stddev_types):
        """
        See :meth:`superclass method
        <.base.GroundShakingIntensityModel.get_mean_and_stddevs>`
        for spec of input and result values.
        """

        # Distance term
        R = np.sqrt(dists.rjb ** 2 + 11.29 ** 2)

        # Magnitude term
        M = rup.mag - 6

        # Site term only distinguishes between lava and ash;
        # since ash sites have Vs30 in the range 60-200m/s,
        # we use this upper value as class separator
        S = np.zeros(R.shape)
        S[sites.vs30 <= 200] = 1

        # Mean ground motion (log10)
        mean = (0.518 + 0.387*M - np.log10(R) - 0.00256*R + 0.335*S)

        # Converting to natural log
        mean /= np.log10(np.e)

        # Check for standard deviation type
        assert all(stddev_type in self.DEFINED_FOR_STANDARD_DEVIATION_TYPES
                   for stddev_type in stddev_types)

        # Constant (total) standard deviation
        stddevs = [0.237/np.log10(np.e) + np.zeros(R.shape)]

        return mean, stddevs

class MunsonThurber1997Vector(MunsonThurber1997):
    """
    Modification of the original base class to correct mean ground motion
    to geometric mean of horizontal components (Beyer and Bommer, 2006)
    """

    #: Supported intensity measure component is geometric mean of horizontal
    #: :attr:`~openquake.hazardlib.const.IMC.VECTORIAL`,
    DEFINED_FOR_INTENSITY_MEASURE_COMPONENT = const.IMC.VECTORIAL

    def get_mean_and_stddevs(self, sites, rup, dists, imt, stddev_types):
        mean, stddevs = super().get_mean_and_stddevs(sites, rup, dists,
                                                     imt, stddev_types)

        # Conversion to geometric mean of horizontal components
        # using the coefficient in Beyer and Bommer, 2006
        mean += np.log(1.1)

        return mean, stddevs


# Licensed under a 3-clause BSD style license - see LICENSE.rst
from copy import deepcopy
import numpy as np
import astropy.io.fits as fits
import astropy.units as u
from gammapy.irf import EDispKernel
from gammapy.maps import Map, MapCoord, WcsGeom
from gammapy.utils.random import InverseCDFSampler, get_random_state

__all__ = ["make_edisp_map", "EDispMap"]

def make_edisp_map(edisp, pointing, geom, exposure_map=None):
    """Make a edisp map for a single observation

    Expected axes : migra and true energy in this specific order
    The name of the migra MapAxis is expected to be 'migra'

    Parameters
    ----------
    edisp : `~gammapy.irf.EnergyDispersion2D`
        the 2D Energy Dispersion IRF
    pointing : `~astropy.coordinates.SkyCoord`
        the pointing direction
    geom : `~gammapy.maps.Geom`
        the map geom to be used. It provides the target geometry.
        rad and true energy axes should be given in this specific order.
    exposure_map : `~gammapy.maps.Map`, optional
        the associated exposure map.
        default is None

    Returns
    -------
    edispmap : `~gammapy.cube.EDispMap`
        the resulting EDisp map
    """
    energy_axis = geom.get_axis_by_name("energy")
    energy = energy_axis.center

    migra_axis = geom.get_axis_by_name("migra")
    migra = migra_axis.center

    # Compute separations with pointing position
    offset = geom.separation(pointing)

    # Compute EDisp values
    edisp_values = edisp.data.evaluate(
        offset=offset,
        e_true=energy[:, np.newaxis, np.newaxis, np.newaxis],
        migra=migra[:, np.newaxis, np.newaxis],
    )

    # Create Map and fill relevant entries
    data = edisp_values.to_value("")
    edispmap = Map.from_geom(geom, data=data, unit="")
    return EDispMap(edispmap, exposure_map)

class EDispMap:
    """Energy dispersion map.

    Parameters
    ----------
    edisp_map : `~gammapy.maps.Map`
        the input Energy Dispersion Map. Should be a Map with 2 non spatial axes.
        migra and true energy axes should be given in this specific order.
    exposure_map : `~gammapy.maps.Map`, optional
        Associated exposure map. Needs to have a consistent map geometry.

    Examples
    --------
    ::

        import numpy as np
        from astropy import units as u
        from astropy.coordinates import SkyCoord
        from gammapy.maps import WcsGeom, MapAxis
        from gammapy.irf import EnergyDispersion2D, EffectiveAreaTable2D
        from gammapy.cube import make_edisp_map, make_map_exposure_true_energy

        # Define energy dispersion map geometry
        energy_axis = MapAxis.from_edges(np.logspace(-1, 1, 4), unit="TeV", name="energy")
        migra_axis = MapAxis.from_edges(np.linspace(0, 3, 100), name="migra")
        pointing = SkyCoord(0, 0, unit="deg")
        max_offset = 4 * u.deg
        geom = WcsGeom.create(
            binsz=0.25 * u.deg,
            width=10 * u.deg,
            skydir=pointing,
            axes=[migra_axis, energy_axis],
        )

        # Extract EnergyDispersion2D from CTA 1DC IRF
        filename = "$GAMMAPY_DATA/cta-1dc/caldb/data/cta/1dc/bcf/South_z20_50h/irf_file.fits"
        edisp2D = EnergyDispersion2D.read(filename, hdu="ENERGY DISPERSION")
        aeff2d = EffectiveAreaTable2D.read(filename, hdu="EFFECTIVE AREA")

        # Create the exposure map
        exposure_geom = geom.to_image().to_cube([energy_axis])
        exposure_map = make_map_exposure_true_energy(pointing, "1 h", aeff2d, exposure_geom)

        # create the EDispMap for the specified pointing
        edisp_map = make_edisp_map(edisp2D, pointing, geom, max_offset, exposure_map)

        # Get an Energy Dispersion (1D) at any position in the image
        pos = SkyCoord(2.0, 2.5, unit="deg")
        e_reco = np.logspace(-1.0, 1.0, 10) * u.TeV
        edisp = edisp_map.get_edisp_kernel(pos=pos, e_reco=e_reco)

        # Write map to disk
        edisp_map.write("edisp_map.fits")
    """

    def __init__(self, edisp_map, exposure_map):
        if edisp_map.geom.axes[1].name.upper() != "ENERGY":
            raise ValueError("Incorrect energy axis position in input Map")

        if edisp_map.geom.axes[0].name.upper() != "MIGRA":
            raise ValueError("Incorrect migra axis position in input Map")

        self.edisp_map = edisp_map
        self.exposure_map = exposure_map

    @classmethod
    def from_hdulist(
        cls,
        hdulist,
        edisp_hdu="EDISPMAP",
        edisp_hdubands="BANDSEDISP",
        exposure_hdu="EXPMAP",
        exposure_hdubands="BANDSEXP",
    ):
        """Convert to `~astropy.io.fits.HDUList`.

        Parameters
        ----------
        edisp_hdu : str
            Name or index of the HDU with the edisp_map data.
        edisp_hdubands : str
            Name or index of the HDU with the edisp_map BANDS table.
        exposure_hdu : str
            Name or index of the HDU with the exposure_map data.
        exposure_hdubands : str
            Name or index of the HDU with the exposure_map BANDS table.
        """
        edisp_map = Map.from_hdulist(hdulist, edisp_hdu, edisp_hdubands, "auto")
        if exposure_hdu in hdulist:
            exposure_map = Map.from_hdulist(
                hdulist, exposure_hdu, exposure_hdubands, "auto"
            )
        else:
            exposure_map = None

        return cls(edisp_map, exposure_map)

    @classmethod
    def read(cls, filename, **kwargs):
        """Read an edisp_map from file and create an EDispMap object"""
        with fits.open(filename, memmap=False) as hdulist:
            return cls.from_hdulist(hdulist, **kwargs)

    def to_hdulist(
        self,
        edisp_hdu="EDISPMAP",
        edisp_hdubands="BANDSEDISP",
        exposure_hdu="EXPMAP",
        exposure_hdubands="BANDSEXP",
    ):
        """Convert to `~astropy.io.fits.HDUList`.

        Parameters
        ----------
        edisp_hdu : str
            Name or index of the HDU with the edisp_map data.
        edisp_hdubands : str
            Name or index of the HDU with the edisp_map BANDS table.
        exposure_hdu : str
            Name or index of the HDU with the exposure_map data.
        exposure_hdubands : str
            Name or index of the HDU with the exposure_map BANDS table.

        Returns
        -------
        hdu_list : `~astropy.io.fits.HDUList`
        """
        hdulist = self.edisp_map.to_hdulist(hdu=edisp_hdu, hdu_bands=edisp_hdubands)
        if self.exposure_map is not None:
            new_hdulist = self.exposure_map.to_hdulist(
                hdu=exposure_hdu, hdu_bands=exposure_hdubands
            )
            hdulist.extend(new_hdulist[1:])
        return hdulist

    def write(self, filename, overwrite=False, **kwargs):
        """Write to fits"""
        hdulist = self.to_hdulist(**kwargs)
        hdulist.writeto(filename, overwrite=overwrite)

    def get_edisp_kernel(self, position, e_reco, migra_step=5e-3):
        """Get energy dispersion at a given position.

        Parameters
        ----------
        position : `~astropy.coordinates.SkyCoord`
            the target position. Should be a single coordinates
        e_reco : `~astropy.units.Quantity`
            Reconstructed energy axis binning
        migra_step : float
            Integration step in migration

        Returns
        -------
        edisp : `~gammapy.irf.EnergyDispersion`
            the energy dispersion (i.e. rmf object)
        """
        # TODO: reduce code duplication with EnergyDispersion2D.get_response
        if position.size != 1:
            raise ValueError(
                "EnergyDispersion can be extracted at one single position only."
            )

        # axes ordering fixed. Could be changed.
        pix_ener = np.arange(self.edisp_map.geom.axes[1].nbin)

        # Define a vector of migration with mig_step step
        mrec_min = self.edisp_map.geom.axes[0].edges[0]
        mrec_max = self.edisp_map.geom.axes[0].edges[-1]
        mig_array = np.arange(mrec_min, mrec_max, migra_step)
        pix_migra = (mig_array - mrec_min) / mrec_max * self.edisp_map.geom.axes[0].nbin

        # Convert position to pixels
        pix_lon, pix_lat = self.edisp_map.geom.to_image().coord_to_pix(position)

        # Build the pixels tuple
        pix = np.meshgrid(pix_lon, pix_lat, pix_migra, pix_ener)
        # Interpolate in the EDisp map. Squeeze to remove dimensions of length 1
        edisp_values = self.edisp_map.interp_by_pix(pix) * u.Unit(self.edisp_map.unit)
        edisp_values = np.squeeze(edisp_values, axis=(0, 1))

        e_trues = self.edisp_map.geom.axes[1].center
        data = []

        for i, e_true in enumerate(e_trues):
            # We now perform integration over migra
            # The code is adapted from `~gammapy.EnergyDispersion2D.get_response`

            # migration value of e_reco bounds
            migra_e_reco = e_reco / e_true

            # Compute normalized cumulative sum to prepare integration
            tmp = np.nan_to_num(
                np.cumsum(edisp_values[:, i]) / np.sum(edisp_values[:, i])
            )

            # Determine positions (bin indices) of e_reco bounds in migration array
            pos_mig = np.digitize(migra_e_reco, mig_array) - 1
            # We ensure that no negative values are found
            pos_mig = np.maximum(pos_mig, 0)

            # We compute the difference between 2 successive bounds in e_reco
            # to get integral over reco energy bin
            integral = np.diff(tmp[pos_mig])

            data.append(integral)

        data = np.asarray(data)
        # EnergyDispersion uses edges of true energy bins
        e_true_edges = self.edisp_map.geom.axes[1].edges

        e_lo, e_hi = e_true_edges[:-1], e_true_edges[1:]
        ereco_lo, ereco_hi = (e_reco[:-1], e_reco[1:])

        return EDispKernel(
            e_true_lo=e_lo,
            e_true_hi=e_hi,
            e_reco_lo=ereco_lo,
            e_reco_hi=ereco_hi,
            data=data,
        )

    def stack(self, other, weights=None):
        """Stack EDispMap with another one in place.

        Parameters
        ----------
        other : `~gammapy.cube.EDispMap`
            Energy dispersion map to be stacked with this one.

        """
        if self.exposure_map is None or other.exposure_map is None:
            raise ValueError("Missing exposure map for PSFMap.stack")

        cutout_info = other.edisp_map.geom.cutout_info

        if cutout_info is not None:
            slices = cutout_info["parent-slices"]
            parent_slices = Ellipsis, slices[0], slices[1]
        else:
            parent_slices = None

        self.edisp_map.data[parent_slices] *= self.exposure_map.data[parent_slices]
        self.edisp_map.stack(other.edisp_map * other.exposure_map.data, weights=weights)

        # stack exposure map
        self.exposure_map.stack(other.exposure_map, weights=weights)

        with np.errstate(invalid="ignore"):
            self.edisp_map.data[parent_slices] /= self.exposure_map.data[parent_slices]
            self.edisp_map.data = np.nan_to_num(self.edisp_map.data)

    def copy(self):
        """Copy EDispMap"""
        return deepcopy(self)

    @classmethod
    def from_geom(cls, geom):
        """Create edisp map from geom.

        By default a diagonal edisp matrix is created.

        Parameters
        ----------
        geom : `Geom`
            Edisp map geometry.

        Returns
        -------
        edisp_map : `EDispMap`
            Energy dispersion map.
        """
        geom_exposure_edisp = geom.squash(axis="migra")
        exposure_edisp = Map.from_geom(geom_exposure_edisp, unit="m2 s")

        migra_axis = geom.get_axis_by_name("migra")
        edisp_map = Map.from_geom(geom, unit="")
        loc = migra_axis.edges.searchsorted(1.0)
        edisp_map.data[:, loc, :, :] = 1.0
        return cls(edisp_map, exposure_edisp)

    def sample_coord(self, map_coord, random_state=0):
        """Apply the energy dispersion corrections on the coordinates of a set of simulated events.

        Parameters
        ----------
        map_coord : `~gammapy.maps.MapCoord` object.
            Sequence of coordinates and energies of sampled events.
        random_state : {int, 'random-seed', 'global-rng', `~numpy.random.RandomState`}
            Defines random number generator initialisation.
            Passed to `~gammapy.utils.random.get_random_state`.

        Returns
        -------
        `~gammapy.maps.MapCoord`.
            Sequence of Edisp-corrected coordinates of the input map_coord map.
        """
        random_state = get_random_state(random_state)
        migra_axis = self.edisp_map.geom.get_axis_by_name("migra")

        coord = {
            "skycoord": map_coord.skycoord.reshape(-1, 1),
            "energy": map_coord["energy"].reshape(-1, 1),
            "migra": migra_axis.center,
        }

        pdf_edisp = self.edisp_map.interp_by_coord(coord)

        sample_edisp = InverseCDFSampler(pdf_edisp, axis=1, random_state=random_state)
        pix_edisp = sample_edisp.sample_axis()
        migra = migra_axis.pix_to_coord(pix_edisp)

        energy_reco = map_coord["energy"] * migra

        return MapCoord.create({"skycoord": map_coord.skycoord, "energy": energy_reco})

    @classmethod
    def from_diagonal_response(cls, energy_axis_true, migra_axis=None):
        """Create an allsky EDisp map with diagonal response.

        Parameters
        ----------
        energy_axis_true : `MapAxis`
            True energy axis
        migra_axis : `MapAxis`
            Migra axis

        Returns
        -------
        edisp_map : `EDispMap`
            Energy dispersion map.
        """
        from .fit import MIGRA_AXIS_DEFAULT

        migra_axis = migra_axis or MIGRA_AXIS_DEFAULT

        geom = WcsGeom.create(
            npix=(2, 1), proj="CAR", binsz=180, axes=[migra_axis, energy_axis_true]
        )

        return cls.from_geom(geom)

    def cutout(self, position, width, mode="trim"):
        """Cutout edisp map.

        Parameters
        ----------
        position : `~astropy.coordinates.SkyCoord`
            Center position of the cutout region.
        width : tuple of `~astropy.coordinates.Angle`
            Angular sizes of the region in (lon, lat) in that specific order.
            If only one value is passed, a square region is extracted.
        mode : {'trim', 'partial', 'strict'}
            Mode option for Cutout2D, for details see `~astropy.nddata.utils.Cutout2D`.

        Returns
        -------
        cutout : `EdispMap`
            Cutout edisp map.
        """
        edisp_map = self.edisp_map.cutout(position, width, mode)
        exposure_map = self.exposure_map.cutout(position, width, mode)
        return self.__class__(edisp_map=edisp_map, exposure_map=exposure_map)

import numpy,import astropy,from astropy