"""
Functions to support rapid interactive modification of Gurobi models.

For reference on Gurobi objects such as Models, Variables, and Constraints, see
http://www.gurobi.com/documentation/7.0/refman/py_python_api_overview.html.
"""
import csv
import json

try:
    import gurobipy as gp
except ImportError:
    raise ImportError("gurobipy not installed. Please see {0} to download".format(
        "https://www.gurobi.com/documentation/6.5/quickstart_mac/the_gurobi_python_interfac.html"))

# Assuming that constraints are of the form:
# constraintName(index1,index2,...,indexN).
# Asuming that variables are of the form:
# variableName[index1,index2,...,indexN]
CON_BRACKET_L = "("
CON_BRACKET_R = ")"
VAR_BRACKET_L = "["
VAR_BRACKET_R = "]"

# 13 July 2016 - Need to sort out capitalization here for attributes

# Attributes of a Gurobi variable
VAR_ATTRS = ["LB", "UB", "Obj", "VType", "VarName", "X", "Xn", "RC",
             "BarX", "Start", "VarHintVal", "VarHintPri", "BranchPriority",
             "VBasis", "PStart", "IISLB", "IISUB", "PWLObjCvx",
             "SAObjLow", "SAObjUp", "SALBLow", "SALBUp",
             "SAUBLow", "SAUBUp", "UnbdRay"]

# Attributes of a Gurobi constraint
CON_ATTRS = ["Sense", "RHS", "ConstrName", "Pi", "Slack",
             "CBasis", "DStart", "Lazy", "IISConstr",
             "SARHSLow", "SARHSUp", "FarkasDual"]

def read_model(filename):
    """
    Read a model using gurobipy.
    """

    m = gp.read(filename)

    return m

def reoptimize(m):
    """
    Update, reset, and optimize
    a model.
    """

    m.update()
    m.reset()
    m.optimize()

def get_variable_attrs():
    """
    Return a list of variable attributes.

    Details of attributes found at the Gurobi
    website:
    http://www.gurobi.com/documentation/6.5/refman/attributes.html
    """

    return VAR_ATTRS

def get_constraint_attrs():
    """
    Return a list of constraint attributes.

    Details of attributes found at the Gurobi
    website:
    http://www.gurobi.com/documentation/6.5/refman/attributes.html
    """

    return CON_ATTRS

def list_constraints(model):
    """
    Print to screen the constraint sets in the model.
    Show the name of each constraint set along with the
    number of constraints in that set.

    A constraint set is composed of all constraints
    sharing the same string identifier before the indices:
    A(2,3,4) and A(1,2,3) are in the same constraint set, A;
    A(2,3,4) and B(2,3,4) are in constraint sets A and B, respectively
    """

    sets = {}
    constraints = model.getConstrs()

    # Assuming constraint set name separated from indicies by
    for c in constraints:
        name = c.constrName
        split_name = name.split(CON_BRACKET_L)
        set_name = split_name[0]

        if set_name not in sets:
            sets[set_name] = 1
        else:
            sets[set_name] += 1

    print "Constraint set, Number of constraints"
    print "\n".join(["{0}, {1}".format(name, number) for name, number
                            in sorted(sets.items())])

def list_variables(model):
    """
    Print to screen the variable sets in the model.
    Show the name of each variable set along with the
    number of variables in that set.

    A variable set is composed of all variables
    sharing the same string identifier before the indices:
    A[2,3,4] and A[1,2,3] are in the same variable set, A;
    A[2,3,4] and B[2,3,4] are in variable sets A and B, respectively
    """

    sets = {}
    variables = model.getVars()

    # Assuming constraint set name separated from indicies by
    for v in variables:
        name = v.varName
        split_name = name.split(VAR_BRACKET_L)
        set_name = split_name[0]

        if set_name not in sets:
            sets[set_name] = 1
        else:
            sets[set_name] += 1

    print "Variable set, Number of variables"
    print "\n".join(["{0}, {1}".format(name, number) for name, number
                            in sorted(sets.items())])

def get_variables(model, name="", approx=False, filter_values={}, exclude=False):
    """
    Return a list of variables from the model
    selected by variable set name.

    A variable set is composed of all variables
    sharing the same string identifier before the indices:
    A[2,3,4] and A[1,2,3] are in the same variable set, A;
    A[2,3,4] and B[2,3,4] are in varaible sets A and B, respectively

    PyGurobi by default assumes that *variable names* are separated
    from indices by square brackets "[" and "]",
    For example, variables look like x[i,j] - "x" in the variable set name,
    and "i" and "j" and the variable's index values.
    See the source code for more details.
    """

    variables = []
    if not name:
        variables = model.getVars()

    if not approx:
        variables = [v for v in model.getVars()
                if v.varName.split(VAR_BRACKET_L)[0] == name]
    else:
        variables = [v for v in model.getVars()
                if name in v.varName.split(VAR_BRACKET_L)[0]]

    if filter_values:
        variables = filter_variables(variables, filter_values,
                exclude=exclude)

    return variables

def check_attr(attr, attributes):
    """
    Check if the attr string case-insensitively corresponds to a
    Gurobi attribute.
    """
    for a in attributes:

        if attr == a:
            return True
        if attr.lower() == a.lower():
            return True

    return False

def check_variable_attr(attr):
    """
    Check if a string corresponds to a variable attribute.

    Case-insensitive.
    """

    var_attrs = get_variable_attrs()

    return check_attr(attr, var_attrs)

def check_constraint_attr(attr):
    """
    Check if a string corresponds to a constraint attribute.

    Attributes are case-insensitive.
    """

    con_attrs = get_constraint_attrs()

    return check_attr(attr, con_attrs)

def get_variables_attr(attr, model="", name="", variables=""):
    """
    Return a dictionary of variables names and their
    corresponding attribute value.

    Specifiy either model and name parameters or supply a list of variables

    """

    if not attr:
        raise AttributeError("No attributes specified")

    if not check_variable_attr(attr):
        raise AttributeError("{0}\n{1}\n{2}".format(
        "Attribute: {0} not a variable attribute.".format(attr),
        "Get list of all variables attributes with the",
        "get_variable_attrs() method."))

    # Make a list of attributes at the top and check against
    # them to make sure that the specified attribute belongs.

    if not model and not variables:
        raise ValueError("No model or variable list given")

    variables = variables_check(model, name, variables)

    return {v.varName: getattr(v, attr) for v in variables}

def print_variables_attr(attr, model="", name="", variables=""):
    """
    Print to screen a dictionary of variables names and their
    corresponding attribute value.

    Specifiy either model and name parameters or supply a list of variables

    """

    var_dict = get_variables_attr(attr, model=model,
                                 name=name, variables=variables)

    print "\n".join(["{0}, {1}".format(v, k) for v, k in
                    sorted(var_dict.items())])

def set_variables_attr(attr, val, model="", name="", variables=""):
    """
    Set an attribute of a model variable set.

    Specifiy either model and name parameters or supply a list of variables
    """
    if not attr or not val:
        raise AttributeError("No attribute or value specified")
        return

    if not check_variable_attr(attr):
        raise AttributeError("{0}\n{1}\n{2}".format(
        "Attribute: {0} not a variable attribute.".format(attr),
        "Get list of all variables attributes with the",
        "get_variable_attrs() method."))

    if not model and not variables:
        raise ValueError("No model or variables specified")

    variables = variables_check(model, name, variables)

    for v in variables:
        setattr(v, attr, val)

def zero_all_objective_coeffs(model):
    """
    Set all objective coefficients in a model to zero.
    """

    if not model:
        raise ValueError("No model given")

    for v in model.getVars():
        v.Obj = 0

def set_variables_bounds(lb="", ub="", model="", name="", variables=""):
    """
    Set the lower bound and/or upper bound for a variables set.

    Specifiy either model and name parameters or supply a list of variables
    """

    if lb:
        set_variables_attr("lb", val=lb, model=model,
                          name=name, variables=variables)

    if ub:
        set_variables_attr("ub", val=ub, model=model,
                          name=name, variables=variables)

def remove_variables_from_model(model, name="", variables=""):
    """
    Remove the given variables from the model.

    Specifiy either model and name parameters or supply a list of constraints
    """

    if not model and not variables:
        raise ValueError("No model or variables given")

    if not model:
        raise ValueError("No model given")

    variables = variables_check(model, name, variables)

    for v in variables:
        model.remove(v)

def variables_check(model, name, variables):
    """
    Return the appropriate
    variables based on the information supplied.
    """

    if variables:
        return variables

    if model and name:
        variables = get_variables(model, name)
    if model and not name:
        variables = model.getVars()

    if not variables:
        print "No variables found for\nmodel: {0},\nname: {1}".format(
               model, name)

    return variables

def get_variable_index_value(variable, index):
    """
    Return the value of the given index
    for a given variable.

    Variable names are assumed to be given
    as A[a,c,d, ....,f]
    """

    value = variable.varName.split(",")[index].strip()
    if VAR_BRACKET_R in value:
        value = value[:-1]
    elif VAR_BRACKET_L in value:
        value = value.split(VAR_BRACKET_L)[1]

    # Not expecting many variable index values to
    # to be floats
    if value.isdigit:
        try:
            value = int(value)
        except ValueError:
            pass

    return value

def get_linexp_from_variables(variables):
    """
    Return a linear expression from the supplied list
    of variables.

    """

    linexp = gp.LinExpr()
    for v in variables:
        linexp += v

    return linexp

def sum_variables_by_index(index, model="", name="", variables=""):
    """
    Return a dictionary mapping index values to the sum
    of the solution values of all matching variables.

    Specifiy either model and name parameters or supply a list of variables
    """

    var_dict = get_variables_by_index(index, model=model, name=name,
                                      variables=variables)
    if not var_dict:
        raise ValueError("No variables found".format(index))

    new_dict = {index_name: sum([v.X for v in index_vars])
                for index_name, index_vars in
                sorted(var_dict.items())}

    return new_dict

def print_dict(dictionary):
    """
    Print a dictionary to screen.
    """

    print "\n".join(["{0}, {1}".format(index_name, index_value)
                     for index_name, index_value in
                     sorted(dictionary.items())])

def print_variables_sum_by_index(index, model="", name="", variables=""):
    """
    Print a dictionary of variables, summed by index.
    """

    var_dict = sum_variables_by_index(index, model=model,
                                         name=name, variables=variables)
    print_dict(var_dict)

def get_variables_by_index(index, model="", name="", variables=""):
    """
    Return a dictionary mapping index values to lists of
    matching variables.

    Specifiy either model and name parameters or supply a list of variables

    """

    if index != 0 and not index:
        raise IndexError("No index given")

    if not model and not variables:
        raise ValueError("No model or variables given")

    if not (name and model) and not variables:
        raise ValueError("No variables specified")

    variables = variables_check(model, name, variables)

    var_dict = {}

    for v in variables:

        value = get_variable_index_value(v, index)
        if value not in var_dict:
            var_dict[value] = [v]
        else:
            var_dict[value].append(v)

    return var_dict

def filter_variables(variables, filter_values, exclude=False):
    """
    Return a new list of variables that match the filter values
    from the given variables list.

    """

    if not variables:
        raise ValueError("variables not given")

    if not filter_values:
        raise ValueError("Dictionary of filter values not given")

    new_vars = []
    for v in variables:
        add = True
        for index, value in filter_values.iteritems():
            key = get_variable_index_value(v, index)
            if key != value:
                add = False
                break
        if add:
            new_vars.append(v)

    if exclude:
        new_vars = [v for v in (set(variables)-set(new_vars))]

    return new_vars

def get_variables_by_index_values(model, name, index_values, exclude=False):

    variables = get_variables(model, name, index_values, exclude)

    return variables

def get_variables_by_two_indices(index1, index2, model="", name="", variables=""):
    """
    Return a dictionary of variables mapping index1 values
    to dictionaries mapping
    index2 values to matching variables.

    Specifiy either model and name parameters or supply a list of variables
    """

    two_indices_dict = {}
    index1_dict = get_variables_by_index(index1, model=model, name=name,
                                      variables=variables)
    for key, value in index1_dict.iteritems():
        two_indices_dict[key] = get_variables_by_index(index2, variables=value)

    return two_indices_dict

def print_variables(variables):
    """
    Print a list of variables to look good.
    """

    print "\n".join([v.varName for v in variables])

def sum_variables_by_two_indices(index1, index2, model="", name="", variables=""):
    """
    Return a dictionary mapping index1 values
    to dictionaries of the given variables summed over index2.

    """

    two_indices_dict = get_variables_by_two_indices(index1, index2,
                                        model=model, name=name, variables=variables)
    if not two_indices_dict:
        raise ValueError("Inputs did not match with model variables")

    new_dict = {}
    for key, var_dict in two_indices_dict.iteritems():
        new_dict[key] = {index_name: sum([v.X for v in index_vars])
                for index_name, index_vars in
                sorted(var_dict.items())}

    return new_dict

def print_two_indices_dict(indices_dict):
    """
    Print to screen a two level nested dictionary.
    """

    for key, value in indices_dict.iteritems():
        print "\n{0}".format(key)
        print_dict(value)

def get_linexp_by_index(index, model="", name="", variables=""):
    """
    Return a dictionary of index values to Gurobi linear expressions
    corresponding to the summation of variables that match the index
    value for the given index number.

    Specifiy either model and name parameters or supply a list of variables.
    """

    linexps = {}

    variables = variables_check(model, name, variables)

    for v in variables:

        value = get_variable_index_value(v, index)

        if value not in linexps:
            linexps[value] = gp.LinExpr(v)
        else:
            linexps[value] += v

    return linexps

def print_constraints(constraints):
    """
    Print constraints in an aesthetically pleasing way.
    """

    print "\n".join([c.constrName for c in constraints])

def get_constraints_multiple(model, names_list, approx=False):
    """
    Return a list of constraints given by the constraint
    set names in names_list.
    """

    cons_list = []
    for name in names_list:
        cons_list.extend(get_constraints(model, name, approx))

    return cons_list

def filter_constraints(constraints, filter_values, exclude=False):
    """
    Return a new list of constraints that match the filter values from
    the given constraints list.
    """

    if not constraints:
        raise ValueError("constraints not given")

    if not filter_values:
        raise ValueError("Dictionary of filter values not given")

    new_cons = []
    for c in constraints:
        add = True
        for index, value in filter_values.iteritems():
            key = get_constraint_index_value(c, index)
            try:
                key.replace('"', "")
            except AttributeError:
                pass
            if key != value:
                add = False
                break
        if add:
            new_cons.append(c)

    if exclude:
        # May want to add sorting by varName here
        new_cons = [c for c in (set(constraints)-set(new_cons))]

    return new_cons

def get_constraints(model, name="", approx=False, filter_values={},
        exclude=False):
    """
    Return a list of constraints from the model
    selected by constraint set name.

    A constraint set is composed of all constraints
    sharing the same string identifier before the indices:
    A(2,3,4) and A(1,2,3) are in the same constraint set, A;
    A(2,3,4) and B(2,3,4) are in constraint sets A and B, respectively

    PyGurobi by default assumes that constraint set names are
    separated from indices by round brackets
    "(" and ")". For example, constraints look like env(r,t) - where "env"
    in the constraint set name
    and "r" and "t" are the index values. See the source for more details.
    """

    if not name:
        return model.getConstrs()

    constraints = []
    if not approx:
        constraints = [c for c in model.getConstrs()
                if c.constrName.split(CON_BRACKET_L)[0] == name]
    else:
        constraints = [c for c in model.getConstrs()
                if name in c.constrName.split(CON_BRACKET_L)[0]]

    if filter_values:
        constraints = filter_constraints(constraints, filter_values, exclude)

    return constraints

def constraints_check(model, name, constraints):
    """
    Check to see whether the user specified a list
    of constraints or expects them to be retrieved
    from the model.
    """

    if constraints:
        return constraints

    if model and name:
        constraints = get_constraints(model, name)
    elif model and not name:
        constraints = model.getConstrs()

    return constraints

def get_constraints_attr(attr, model="", name="", constraints=""):
    """
    Return a dictionary of constraint names and their
    corresponding attribute value.

    Specifiy either model and name parameters or supply a list of constraints
    """

    if not attr:
        raise AttributeError("No attributes specified")

    if not check_constraint_attr(attr):
        raise AttributeError("{0}\n{1}\n{2}".format(
        "Attribute: {0} not a constraint attribute.".format(attr),
        "Get list of all variables attributes with the",
        "get_constraint_attrs() method."))

    # Check if the attr supplied is not a viable model attribute
    if not model and not constraints:
        raise ValueError("No model or constraint list given")

    constraints = constraints_check(model, name, constraints)

    return {c.constrName: getattr(c, attr) for c in constraints}

def print_constraints_attr(attr, model="", name="", constraints=""):
    """
    Print to screen a list of constraint attribute values
    given by the constraints specified in the names parameter.

    Specifiy either model and name parameters or supply a list of constraints

    """

    constraints = get_constraints_attr(attr, model=model,
                                      name=name, constraints=constraints)

    print "\n".join(["{0}, {1}".format(c, k)
                     for c, k in sorted(constraints.items())])

def set_constraints_attr(attr, val, model="", name="", constraints=""):
    """
    Set an attribute of a model constraint set.

    Specifiy either model and name parameters or supply a list of constraints
    """

    if not attr or not val:
        raise AttributeError("No attribute or value specified")

    if not check_constraint_attr(attr):
        raise AttributeError("{0}\n{1}\n{2}".format(
        "Attribute: {0} not a variable attribute.".format(attr),
        "Get list of all variables attributes with the",
        "get_variable_attrs() method."))

    if not model and not constraints:
        raise ValueError("No model or constraints specified")

    constraints = constraints_check(model, name, constraints)

    for c in constraints:
        setattr(c, attr, val)

def set_constraints_rhs_as_percent(percent, model="", name="", constraints=""):
    """
    Set the right hand side (rhs) of a constraint set as a percentage of its current rhs.

    Specifiy either model and name parameters or supply a list of constraints
    """

    if percent != 0 and not percent:
        print "Error: No percent specified."
        return

    try:
        percent = float(percent)
    except ValueError:
        raise ValueError("Percent must be a number. Percent: {}".format(percent))

    if not model and not constraints:
        raise ValueError("No model or constraints specified.")

    constraints = constraints_check(model, name, constraints)

    for c in constraints:
        cur_rhs = getattr(c, "rhs")
        setattr(c, "rhs", percent*cur_rhs)

def remove_constraints_from_model(model, name="", constraints=""):
    """
    Remove the given constraints from the model.

    Specifiy either model and name parameters or supply a list of constraints
    """

    if not model and not constraints:
        raise ValueError("No model or constraints given")

    if not model:
        raise ValueError("No model given")

    # This is needed for the case where a list of
    # constraints is provided because a model object
    # must be provided
    if not constraints:
        constraints = constraints_check(model, name, constraints)

    for c in constraints:
        model.remove(c)

def get_constraint_index_value(constraint, index):
    """
    Return the value of the given index
    for a given constraint.

    Constraint names are assumed to be given
    as A(a,c,d, ....,f)
    """

    value = constraint.constrName.split(",")[index].strip()
    if CON_BRACKET_R in value:
        value = value[:-1]
    elif CON_BRACKET_L in value:
        value = value.split(CON_BRACKET_L)[1]

    # Not expecting many constraint index values to
    # to be floats
    if value.isdigit:
        try:
            value = int(value)
        except ValueError:
            pass

    return value

def get_constraints_by_index(index, model="", name="", constraints=""):
    """
    Return a dictionary mapping index values to lists of
    constraints having that index value.

    Specifiy either model and name parameters or supply a list of constraints
    """

    if index != 0 and not index:
        raise IndexError("No index given")

    if not model and not constraints:
        raise ValueError("No model or constraints given")

    if not (name and model) and not constraints:
        raise ValueError("No constraints specified")

    constraints = constraints_check(model, name, constraints)

    con_dict = {}

    for c in constraints:

        value = get_constraint_index_value(c, index)
        if value not in con_dict:
            con_dict[value] = [c]
        else:
            con_dict[value].append(c)

    return con_dict

def get_constraints_by_index_values(model, name, index_values, exclude=False):
    """
    Return a list of constraints filtered by index values.

    If exlude is False then return constraints that match the filters.
    If exclude is True than return constraints that do not match the filters.
    """

    constraints = get_constraints(model, name, index_values, exclude)

    return constraints

def get_grb_sense_from_string(sense):
    """
    Return the GRB constraint sense object
    corresponding to the supplied string.

    Convention follows the Gurobi docs:
    https://www.gurobi.com/documentation/6.5/refman/sense.html#attr:Sense
    """

    if sense == "<":
        return gp.GRB.LESS_EQUAL
    elif sense == ">":
        return gp.GRB.GREATER_EQUAL
    elif sense == "=":
        return gp.GRB.EQUAL
    else:
        raise ValueError("Constraint sense is not '<', '>', '='")

def add_constraint_constant(model, variables, constant, sense="<",
                                      con_name=""):
    """
    Add constraint to model that says the sum of
    variables must be equal, less than or equal, or, greater than or equal, a constant.

    """

    if not variables:
        raise ValueError("variables list not provided")

    linexp = get_linexp_from_variables(variables)

    sense = get_grb_sense_from_string(sense)

    if not con_name:
        model.addConstr(linexp, sense, constant)
    else:
        model.addConstr(linexp, sense, constant, con_name)

def check_if_name_a_variable(name, model):
    """
    Check if the supplied name corresponds to
    a variable set name in the given model.
    """

    variables = get_variables(model, name)

    if not variables:
        return False

    return True

def check_if_name_a_constraint(name, model):
    """
    Check if the supplied name corresopnd to
    a constraint set name in the given model.
    """

    constraints = get_constraints(model, name)

    if not constraints:
        return False

    return True

def add_constraint_variables(model, variables1, variables2,
                                 sense="=", con_name=""):
    """
    Add constraint to model that says the sum of
    a list of variables must be equal, less than or equal,
    or greater than or equal, the sum of another list of variables.
    """

    if not variables1 or not variables2:
        ValueError("Variables list not provided")

    linexp1 = get_linexp_from_variables(variables1)
    linexp2 = get_linexp_from_variables(variables2)

    sense = get_grb_sense_from_string(sense)

    if not con_name:
        model.addConstr(linexp1, sense, linexp2)
    else:
        model.addConstr(linexp1, sense, linexp2, con_name)

def graph_by_index(model, variables, index, title="", y_axis="", x_axis=""):
    """
    Display a graph of the variable against the specified index
    using matplotlib.

    Matplotlib must already be installed to use this.
    See: http://matplotlib.org/faq/installing_faq.html
    """
    try:
        import matplotlib.pyplot as plot
    except ImportError:
        raise ImportError("{0}\n{1}".format(
        "Module Matplotlib not found.",
        "Please download and install Matplotlib to use this function."))

    fig = plot.figure()
    ax = fig.add_subplot(111)

    variables_sum = sum_variables_by_index(index, variables=variables)

    keys, values = zip(*variables_sum.items())

    y = range(len(variables_sum))

    if title:
        ax.set_title(title)
    if y_axis:
        ax.set_ylabel(y_axis)
    if x_axis:
        ax.set_xlabel(x_axis)
    ax.bar(y, values)
    #ax.legend(keys)

    plot.show()

def graph_by_two_indices(model, variables, index1, index2, title="",
                                 y_axis="", x_axis=""):
    """
    Display a graph of the variable summed over index2
    given by index1.

    Matplotlib must already be installed to use this.
    See: http://matplotlib.org/faq/installing_faq.html
    """
    try:
        import matplotlib.pyplot as plot
    except ImportError:
        raise ImportError("{0}\n{1}".format(
        "Module Matplotlib not found.",
        "Please download and install Matplotlib to use this function."))

    fig = plot.figure()
    ax = fig.add_subplot(111)

    # We need to do this in reverse order to prepare it for graphing
    variables_sum = sum_variables_by_two_indices(index2, index1,
                                       variables=variables)

    keys, values = zip(*variables_sum.items())

    colours = ["b", "g", "r", "c", "y", "m", "k", "w"]

    y = range(len(values[0]))
    print y

    if title:
        ax.set_title(title)
    if y_axis:
        ax.set_ylabel(y_axis)
    if x_axis:
        ax.set_xlabel(x_axis)
    bars = []

    prev_bars = [0 for bar in y]
    colour_count = 0
    for key, value in variables_sum.iteritems():
        cur_bars = [k[1] for k in sorted(value.items(), key=lambda x: x[0])]
        bars.append(ax.bar(y, cur_bars, bottom=prev_bars,
                                color=colours[colour_count]))
        prev_bars = cur_bars
        colour_count += 1
        if colour_count == len(colours) - 1:
            colour_count = 0
    ax.legend(keys)

    plot.show()

def print_variables_to_csv(file_name, model="", name="", variables=""):
    """
    Print the specified variables to a csv file
    given by the file_name parameter.

    If no variables specified than all model
    variables written.
    """

    if ".csv" not in file_name:
        raise ValueError("Non csv file specified")

    with open(file_name, "wb+") as write_file:
        writer = csv.writer(write_file)

        headers = ["Variable name", "Value"]
        writer.writerow(headers)

        variables = variables_check(model, name, variables)

        # This will put quotes around strings, because the variable
        # names have commas in them.
        writer.writerows([ [v.varName, v.X] for v in variables])

def print_variables_to_csv_by_index(file_name, index,
                                    model="", name="", variables=""):
    """
    Print the sums of variables by the specified index
    to a csv file.

    Default behaviour of the function is to overwrite
    the given file_name.
    """

    if ".csv" not in file_name:
        raise ValueError("Non csv file specified")

    with open(file_name, "wb+") as write_file:
        writer = csv.writer(write_file)

        headers = ["Index", "Value"]
        writer.writerow(headers)

        variables_dict = sum_variables_by_index(index, model=model,
                                            name=name, variables=variables)

        if not variables_dict:
            raise ValueError("No variables found")

        writer.writerows([ [key, value]
                        for key, value in sorted(variables_dict.items())])

def print_variables_to_json_by_index(file_name, index, model="",
                                    name="", variables="", index_alias=""):
    """
    Print the specified variables to a json file given by file_name
    organized by the specified index.

    Formatted for reading into nvD3 applications.

    Default behaviour is to overwrite file if one exists in
    file_name's location.
    """

    if ".json" not in file_name:
        raise ValueError("Non json file specified")

    index_name = index
    if index_alias:
        index_name = index_alias

    var_dict = sum_variables_by_index(index, model=model,
                                      name=name, variables=variables)

    data = {index_name: [{ index_name: var_dict }] }

    json.dump(data, open(file_name, "wb"))


# -*- coding: utf-8 -*-
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

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

import os
import random

from airflow.models import Connection
from airflow.exceptions import AirflowException
from airflow.utils.db import provide_session
from airflow.utils.log.logging_mixin import LoggingMixin

CONN_ENV_PREFIX = 'AIRFLOW_CONN_'

class BaseHook(LoggingMixin):
    """
    Abstract base class for hooks, hooks are meant as an interface to
    interact with external systems. MySqlHook, HiveHook, PigHook return
    object that can handle the connection and interaction to specific
    instances of these systems, and expose consistent methods to interact
    with them.
    """
    def __init__(self, source):
        pass

    @classmethod
    @provide_session
    def _get_connections_from_db(cls, conn_id, session=None):
        db = (
            session.query(Connection)
            .filter(Connection.conn_id == conn_id)
            .all()
        )
        session.expunge_all()
        if not db:
            raise AirflowException(
                "The conn_id `{0}` isn't defined".format(conn_id))
        return db

    @classmethod
    def _get_connection_from_env(cls, conn_id):
        environment_uri = os.environ.get(CONN_ENV_PREFIX + conn_id.upper())
        conn = None
        if environment_uri:
            conn = Connection(conn_id=conn_id, uri=environment_uri)
        return conn

    @classmethod
    def get_connections(cls, conn_id):
        conn = cls._get_connection_from_env(conn_id)
        if conn:
            conns = [conn]
        else:
            conns = cls._get_connections_from_db(conn_id)
        return conns

    @classmethod
    def get_connection(cls, conn_id):
        conn = random.choice(cls.get_connections(conn_id))
        if conn.host:
            log = LoggingMixin().log
            log.info("Using connection to: %s", conn.host)
        return conn

    @classmethod
    def get_hook(cls, conn_id):
        connection = cls.get_connection(conn_id)
        return connection.get_hook()

    def get_conn(self):
        raise NotImplementedError()

    def get_records(self, sql):
        raise NotImplementedError()

    def get_pandas_df(self, sql):
        raise NotImplementedError()

    def run(self, sql):
        raise NotImplementedError()


"""
A deep neural network with or w/o dropout in one file.
"""

import numpy
import theano
import sys
import math
from theano import tensor as T
from theano import shared
from theano.tensor.shared_randomstreams import RandomStreams
from collections import OrderedDict

BATCH_SIZE = 100
STACKSIZE = 69

def relu_f(vec):
    """ Wrapper to quickly change the rectified linear unit function """
    return (vec + abs(vec)) / 2.

def softplus_f(v):
    return T.nnet.softplus(v)

def dropout(rng, x, p=0.5):
    """ Zero-out random values in x with probability p using rng """
    if p > 0. and p < 1.:
        seed = rng.randint(2 ** 30)
        srng = theano.tensor.shared_randomstreams.RandomStreams(seed)
        mask = srng.binomial(n=1, p=1.-p, size=x.shape,
                dtype=theano.config.floatX)
        return x * mask
    return x

def fast_dropout(rng, x):
    """ Multiply activations by N(1,1) """
    seed = rng.randint(2 ** 30)
    srng = RandomStreams(seed)
    mask = srng.normal(size=x.shape, avg=1., dtype=theano.config.floatX)
    return x * mask

def build_shared_zeros(shape, name):
    """ Builds a theano shared variable filled with a zeros numpy array """
    return shared(value=numpy.zeros(shape, dtype=theano.config.floatX),
            name=name, borrow=True)

class Linear(object):
    """ Basic linear transformation layer (W.X + b) """
    def __init__(self, rng, input, n_in, n_out, W=None, b=None, fdrop=False):
        if W is None:
            W_values = numpy.asarray(rng.uniform(
                low=-numpy.sqrt(6. / (n_in + n_out)),
                high=numpy.sqrt(6. / (n_in + n_out)),
                size=(n_in, n_out)), dtype=theano.config.floatX)
            W_values *= 4  # This works for sigmoid activated networks!
            W = theano.shared(value=W_values, name='W', borrow=True)
        if b is None:
            b = build_shared_zeros((n_out,), 'b')
        self.input = input
        self.W = W
        self.b = b
        self.params = [self.W, self.b]
        self.output = T.dot(self.input, self.W) + self.b
        if fdrop:
            self.output = fast_dropout(rng, self.output)

    def __repr__(self):
        return "Linear"

class SigmoidLayer(Linear):
    """ Sigmoid activation layer (sigmoid(W.X + b)) """
    def __init__(self, rng, input, n_in, n_out, W=None, b=None, fdrop=False):
        super(SigmoidLayer, self).__init__(rng, input, n_in, n_out, W, b)
        self.pre_activation = self.output
        if fdrop:
            self.pre_activation = fast_dropout(rng, self.pre_activation)
        self.output = T.nnet.sigmoid(self.pre_activation)

class ReLU(Linear):
    """ Rectified Linear Unit activation layer (max(0, W.X + b)) """
    def __init__(self, rng, input, n_in, n_out, W=None, b=None, fdrop=False):
        if b is None:
            b = build_shared_zeros((n_out,), 'b')
        super(ReLU, self).__init__(rng, input, n_in, n_out, W, b)
        self.pre_activation = self.output
        if fdrop:
            self.pre_activation = fast_dropout(rng, self.pre_activation)
        self.output = relu_f(self.pre_activation)

class SoftPlus(Linear):
    def __init__(self, rng, input, n_in, n_out, W=None, b=None, fdrop=0.):
        if b is None:
            b_values = numpy.zeros((n_out,), dtype=theano.config.floatX)
            b = theano.shared(value=b_values, name='b', borrow=True)
        super(SoftPlus, self).__init__(rng, input, n_in, n_out, W, b)
        self.pre_activation = self.output
        if fdrop:
            self.pre_activation = fast_dropout(rng, self.pre_activation, fdrop)
        self.output = softplus_f(self.pre_activation)

class DatasetMiniBatchIterator(object):
    """ Basic mini-batch iterator """
    def __init__(self, x, y, batch_size=BATCH_SIZE, randomize=False):
        self.x = x
        self.y = y
        self.batch_size = batch_size
        self.randomize = randomize
        from sklearn.utils import check_random_state
        self.rng = check_random_state(42)

    def __iter__(self):
        n_samples = self.x.shape[0]
        if self.randomize:
            for _ in xrange(n_samples / BATCH_SIZE):
                if BATCH_SIZE > 1:
                    i = int(self.rng.rand(1) * ((n_samples+BATCH_SIZE-1) / BATCH_SIZE))
                else:
                    i = int(math.floor(self.rng.rand(1) * n_samples))
                yield (i, self.x[i*self.batch_size:(i+1)*self.batch_size],
                       self.y[i*self.batch_size:(i+1)*self.batch_size])
        else:
            for i in xrange((n_samples + self.batch_size - 1)
                            / self.batch_size):
                yield (self.x[i*self.batch_size:(i+1)*self.batch_size],
                       self.y[i*self.batch_size:(i+1)*self.batch_size])

class LogisticRegression:
    """Multi-class Logistic Regression
    """
    def __init__(self, rng, input, n_in, n_out, W=None, b=None):
        if W != None:
            self.W = W
        else:
            self.W = build_shared_zeros((n_in, n_out), 'W')
        if b != None:
            self.b = b
        else:
            self.b = build_shared_zeros((n_out,), 'b')

        # P(Y|X) = softmax(W.X + b)
        self.p_y_given_x = T.nnet.softmax(T.dot(input, self.W) + self.b)
        self.y_pred = T.argmax(self.p_y_given_x, axis=1)
        self.output = self.y_pred
        self.params = [self.W, self.b]

    def negative_log_likelihood(self, y):
        return -T.mean(T.log(self.p_y_given_x)[T.arange(y.shape[0]), y])

    def negative_log_likelihood_sum(self, y):
        return -T.sum(T.log(self.p_y_given_x)[T.arange(y.shape[0]), y])

    def log_loss(self, y):
        # TODO
        log_y_hat = T.log(self.p_y_given_x)
        #ll = log_y_hat[T.arange(y.shape[0]), y] + log_y_hat[T.arange(y.shape[0]), 1-y]
        #return -T.mean(ll)

    def training_cost(self, y):
        """ Wrapper for standard name """
        return self.negative_log_likelihood_sum(y)
        #return self.log_loss(y) TODO

    def errors(self, y):
        if y.ndim != self.y_pred.ndim:
            raise TypeError("y should have the same shape as self.y_pred",
                ("y", y.type, "y_pred", self.y_pred.type))
        if y.dtype.startswith('int'):
            return T.mean(T.neq(self.y_pred, y))
        else:
            print("!!! y should be of int type")
            return T.mean(T.neq(self.y_pred, numpy.asarray(y, dtype='int')))

class NeuralNet(object):
    """ Neural network (not regularized, without dropout) """
    def __init__(self, numpy_rng, theano_rng=None,
                 n_ins=40*3,
                 layers_types=[Linear, ReLU, ReLU, ReLU, LogisticRegression],
                 layers_sizes=[1024, 1024, 1024, 1024],
                 n_outs=62 * 3,
                 rho=0.95, eps=1.E-6,
                 max_norm=0.,
                 debugprint=False):
        """
        TODO
        """
        self.layers = []
        self.params = []
        self.n_layers = len(layers_types)
        self.layers_types = layers_types
        assert self.n_layers > 0
        self.max_norm = max_norm
        self._rho = rho  # ``momentum'' for adadelta
        self._eps = eps  # epsilon for adadelta
        self._accugrads = []  # for adadelta
        self._accudeltas = []  # for adadelta
        self._old_dxs = []  # for adadelta with Nesterov

        if theano_rng == None:
            theano_rng = RandomStreams(numpy_rng.randint(2 ** 30))

        self.x = T.fmatrix('x')
        self.y = T.ivector('y')

        self.layers_ins = [n_ins] + layers_sizes
        self.layers_outs = layers_sizes + [n_outs]

        layer_input = self.x

        for layer_type, n_in, n_out in zip(layers_types,
                self.layers_ins, self.layers_outs):
            this_layer = layer_type(rng=numpy_rng,
                    input=layer_input, n_in=n_in, n_out=n_out)
            assert hasattr(this_layer, 'output')
            self.params.extend(this_layer.params)
            self._accugrads.extend([build_shared_zeros(t.shape.eval(),
                'accugrad') for t in this_layer.params])
            self._accudeltas.extend([build_shared_zeros(t.shape.eval(),
                'accudelta') for t in this_layer.params])
            self._old_dxs.extend([build_shared_zeros(t.shape.eval(),
                'old_dxs') for t in this_layer.params])

            self.layers.append(this_layer)
            layer_input = this_layer.output

        assert hasattr(self.layers[-1], 'training_cost')
        assert hasattr(self.layers[-1], 'errors')
        # TODO standardize cost
        self.mean_cost = self.layers[-1].negative_log_likelihood(self.y)
        self.cost = self.layers[-1].training_cost(self.y)
        #self.mean_cost = self.layers[-1].training_cost(self.y) # TODO
        if debugprint:
            theano.printing.debugprint(self.cost)

        self.errors = self.layers[-1].errors(self.y)

    def __repr__(self):
        dimensions_layers_str = map(lambda x: "x".join(map(str, x)),
                                    zip(self.layers_ins, self.layers_outs))
        return "_".join(map(lambda x: "_".join((x[0].__name__, x[1])),
                            zip(self.layers_types, dimensions_layers_str)))

    def get_SGD_trainer(self):
        """ Returns a plain SGD minibatch trainer with learning rate as param.
        """
        batch_x = T.fmatrix('batch_x')
        batch_y = T.ivector('batch_y')
        learning_rate = T.fscalar('lr')  # learning rate to use
        # compute the gradients with respect to the model parameters
        # using mean_cost so that the learning rate is not too dependent
        # on the batch size
        gparams = T.grad(self.mean_cost, self.params)

        # compute list of weights updates
        updates = OrderedDict()
        for param, gparam in zip(self.params, gparams):
            if self.max_norm:
                W = param - gparam * learning_rate
                col_norms = W.norm(2, axis=0)
                desired_norms = T.clip(col_norms, 0, self.max_norm)
                updates[param] = W * (desired_norms / (1e-6 + col_norms))
            else:
                updates[param] = param - gparam * learning_rate

        train_fn = theano.function(inputs=[theano.Param(batch_x),
                                           theano.Param(batch_y),
                                           theano.Param(learning_rate)],
                                   outputs=self.mean_cost,
                                   updates=updates,
                                   givens={self.x: batch_x, self.y: batch_y})

        return train_fn

    def get_adagrad_trainer(self):
        """ Returns an Adagrad (Duchi et al. 2010) trainer using a learning rate.
        """
        batch_x = T.fmatrix('batch_x')
        batch_y = T.ivector('batch_y')
        learning_rate = T.fscalar('lr')  # learning rate to use
        # compute the gradients with respect to the model parameters
        gparams = T.grad(self.mean_cost, self.params)

        # compute list of weights updates
        updates = OrderedDict()
        for accugrad, param, gparam in zip(self._accugrads, self.params, gparams):
            # c.f. Algorithm 1 in the Adadelta paper (Zeiler 2012)
            agrad = accugrad + gparam * gparam
            dx = - (learning_rate / T.sqrt(agrad + self._eps)) * gparam
            if self.max_norm:
                W = param + dx
                col_norms = W.norm(2, axis=0)
                desired_norms = T.clip(col_norms, 0, self.max_norm)
                updates[param] = W * (desired_norms / (1e-6 + col_norms))
            else:
                updates[param] = param + dx
            updates[accugrad] = agrad

        train_fn = theano.function(inputs=[theano.Param(batch_x),
            theano.Param(batch_y),
            theano.Param(learning_rate)],
            outputs=self.mean_cost,
            updates=updates,
            givens={self.x: batch_x, self.y: batch_y})

        return train_fn

    def get_adadelta_trainer(self):
        """ Returns an Adadelta (Zeiler 2012) trainer using self._rho and
        self._eps params.
        """
        batch_x = T.fmatrix('batch_x')
        batch_y = T.ivector('batch_y')
        # compute the gradients with respect to the model parameters
        gparams = T.grad(self.mean_cost, self.params)

        # compute list of weights updates
        updates = OrderedDict()
        for accugrad, accudelta, param, gparam in zip(self._accugrads,
                self._accudeltas, self.params, gparams):
            # c.f. Algorithm 1 in the Adadelta paper (Zeiler 2012)
            agrad = self._rho * accugrad + (1 - self._rho) * gparam * gparam
            dx = - T.sqrt((accudelta + self._eps)
                          / (agrad + self._eps)) * gparam
            updates[accudelta] = (self._rho * accudelta
                                  + (1 - self._rho) * dx * dx)
            if self.max_norm:
                W = param + dx
                col_norms = W.norm(2, axis=0)
                desired_norms = T.clip(col_norms, 0, self.max_norm)
                updates[param] = W * (desired_norms / (1e-6 + col_norms))
            else:
                updates[param] = param + dx
            updates[accugrad] = agrad

        train_fn = theano.function(inputs=[theano.Param(batch_x),
                                           theano.Param(batch_y)],
                                   outputs=self.mean_cost,
                                   updates=updates,
                                   givens={self.x: batch_x, self.y: batch_y})

        return train_fn

    def score_classif(self, given_set):
        """ Returns functions to get current classification errors. """
        batch_x = T.fmatrix('batch_x')
        batch_y = T.ivector('batch_y')
        score = theano.function(inputs=[theano.Param(batch_x),
                                        theano.Param(batch_y)],
                                outputs=self.errors,
                                givens={self.x: batch_x, self.y: batch_y})

        def scoref():
            """ returned function that scans the entire set given as input """
            return [score(batch_x, batch_y) for batch_x, batch_y in given_set]

        return scoref

class RegularizedNet(NeuralNet):
    """ Neural net with L1 and L2 regularization """
    def __init__(self, numpy_rng, theano_rng=None,
                 n_ins=100,
                 layers_types=[ReLU, ReLU, ReLU, LogisticRegression],
                 layers_sizes=[1024, 1024, 1024],
                 n_outs=2,
                 rho=0.9, eps=1.E-6,
                 L1_reg=0.,
                 L2_reg=0.,
                 max_norm=0.,
                 debugprint=False):
        """
        TODO
        """
        super(RegularizedNet, self).__init__(numpy_rng, theano_rng, n_ins,
                layers_types, layers_sizes, n_outs, rho, eps, max_norm,
                debugprint)

        L1 = shared(0.)
        for param in self.params:
            L1 += T.sum(abs(param))
        if L1_reg > 0.:
            self.cost = self.cost + L1_reg * L1
        L2 = shared(0.)
        for param in self.params:
            L2 += T.sum(param ** 2)
        if L2_reg > 0.:
            self.cost = self.cost + L2_reg * L2

class DropoutNet(NeuralNet):
    """ Neural net with dropout (see Hinton's et al. paper) """
    def __init__(self, numpy_rng, theano_rng=None,
                 n_ins=40*3,
                 layers_types=[ReLU, ReLU, ReLU, ReLU, LogisticRegression],
                 layers_sizes=[4000, 4000, 4000, 4000],
                 dropout_rates=[0.2, 0.5, 0.5, 0.5, 0.5],
                 n_outs=62 * 3,
                 rho=0.98, eps=1.E-6,
                 max_norm=0.,
                 fast_drop=True,
                 debugprint=False):
        """
        TODO
        """
        super(DropoutNet, self).__init__(numpy_rng, theano_rng, n_ins,
                layers_types, layers_sizes, n_outs, rho, eps, max_norm,
                debugprint)

        self.dropout_rates = dropout_rates
        if fast_drop:
            if dropout_rates[0]:
                dropout_layer_input = fast_dropout(numpy_rng, self.x)
            else:
                dropout_layer_input = self.x
        else:
            dropout_layer_input = dropout(numpy_rng, self.x, p=dropout_rates[0])
        self.dropout_layers = []

        for layer, layer_type, n_in, n_out, dr in zip(self.layers,
                layers_types, self.layers_ins, self.layers_outs,
                dropout_rates[1:] + [0]):  # !!! we do not dropout anything
                                           # from the last layer !!!
            if dr:
                if fast_drop:
                    this_layer = layer_type(rng=numpy_rng,
                            input=dropout_layer_input, n_in=n_in, n_out=n_out,
                            W=layer.W, b=layer.b, fdrop=True)
                else:
                    this_layer = layer_type(rng=numpy_rng,
                            input=dropout_layer_input, n_in=n_in, n_out=n_out,
                            W=layer.W * 1. / (1. - dr),
                            b=layer.b * 1. / (1. - dr))
                    # N.B. dropout with dr==1 does not dropanything!!
                    this_layer.output = dropout(numpy_rng, this_layer.output, dr)
            else:
                this_layer = layer_type(rng=numpy_rng,
                        input=dropout_layer_input, n_in=n_in, n_out=n_out,
                        W=layer.W, b=layer.b)

            assert hasattr(this_layer, 'output')
            self.dropout_layers.append(this_layer)
            dropout_layer_input = this_layer.output

        assert hasattr(self.layers[-1], 'training_cost')
        assert hasattr(self.layers[-1], 'errors')
        # TODO standardize cost
        # these are the dropout costs
        self.mean_cost = self.dropout_layers[-1].negative_log_likelihood(self.y)
        self.cost = self.dropout_layers[-1].training_cost(self.y)

        # these is the non-dropout errors
        self.errors = self.layers[-1].errors(self.y)

    def __repr__(self):
        return super(DropoutNet, self).__repr__() + "\n"\
                + "dropout rates: " + str(self.dropout_rates)

def add_fit_and_score(class_to_chg):
    """ Mutates a class to add the fit() and score() functions to a NeuralNet.
    """
    from types import MethodType
    def fit(self, x_train, y_train, x_dev=None, y_dev=None,
            max_epochs=20, early_stopping=True, split_ratio=0.1, # TODO 100+ epochs
            method='adadelta', verbose=False, plot=False):
        """
        TODO
        """
        import time, copy
        if x_dev == None or y_dev == None:
            from sklearn.cross_validation import train_test_split
            x_train, x_dev, y_train, y_dev = train_test_split(x_train, y_train,
                    test_size=split_ratio, random_state=42)
        if method == 'sgd':
            train_fn = self.get_SGD_trainer()
        elif method == 'adagrad':
            train_fn = self.get_adagrad_trainer()
        elif method == 'adadelta':
            train_fn = self.get_adadelta_trainer()
        elif method == 'adadelta_rprop':
            train_fn = self.get_adadelta_rprop_trainer()
        train_set_iterator = DatasetMiniBatchIterator(x_train, y_train)
        dev_set_iterator = DatasetMiniBatchIterator(x_dev, y_dev)
        train_scoref = self.score_classif(train_set_iterator)
        dev_scoref = self.score_classif(dev_set_iterator)
        best_dev_loss = numpy.inf
        epoch = 0
        # TODO early stopping (not just cross val, also stop training)
        if plot:
            verbose = True
            self._costs = []
            self._train_errors = []
            self._dev_errors = []
            self._updates = []

        while epoch < max_epochs:
            if not verbose:
                sys.stdout.write("\r%0.2f%%" % (epoch * 100./ max_epochs))
                sys.stdout.flush()
            avg_costs = []
            timer = time.time()
            for x, y in train_set_iterator:
                if method == 'sgd' or 'adagrad' in method:
                    avg_cost = train_fn(x, y, lr=1.E-2)
                elif 'adadelta' in method:
                    avg_cost = train_fn(x, y)
                if type(avg_cost) == list:
                    avg_costs.append(avg_cost[0])
                else:
                    avg_costs.append(avg_cost)
            if verbose:
                mean_costs = numpy.mean(avg_costs)
                mean_train_errors = numpy.mean(train_scoref())
                print('  epoch %i took %f seconds' %
                      (epoch, time.time() - timer))
                print('  epoch %i, avg costs %f' %
                      (epoch, mean_costs))
                print('  method %s, epoch %i, training error %f' %
                      (method, epoch, mean_train_errors))
                if plot:
                    self._costs.append(mean_costs)
                    self._train_errors.append(mean_train_errors)
            dev_errors = numpy.mean(dev_scoref())
            if plot:
                self._dev_errors.append(dev_errors)
            if dev_errors < best_dev_loss:
                best_dev_loss = dev_errors
                best_params = copy.deepcopy(self.params)
                if verbose:
                    print('!!!  epoch %i, validation error of best model %f' %
                          (epoch, dev_errors))
            epoch += 1
        if not verbose:
            print("")
        for i, param in enumerate(best_params):
            self.params[i] = param

    def score(self, x, y):
        """ error rates """
        iterator = DatasetMiniBatchIterator(x, y)
        scoref = self.score_classif(iterator)
        return numpy.mean(scoref())

    class_to_chg.fit = MethodType(fit, None, class_to_chg)
    class_to_chg.score = MethodType(score, None, class_to_chg)

if __name__ == "__main__":
    add_fit_and_score(DropoutNet)
    add_fit_and_score(RegularizedNet)

    def nudge_dataset(X, Y):
        """
        This produces a dataset 5 times bigger than the original one,
        by moving the 8x8 images in X around by 1px to left, right, down, up
        """
        from scipy.ndimage import convolve
        direction_vectors = [
            [[0, 1, 0],
             [0, 0, 0],
             [0, 0, 0]],
            [[0, 0, 0],
             [1, 0, 0],
             [0, 0, 0]],
            [[0, 0, 0],
             [0, 0, 1],
             [0, 0, 0]],
            [[0, 0, 0],
             [0, 0, 0],
             [0, 1, 0]]]
        shift = lambda x, w: convolve(x.reshape((8, 8)), mode='constant',
                                      weights=w).ravel()
        X = numpy.concatenate([X] +
                              [numpy.apply_along_axis(shift, 1, X, vector)
                                  for vector in direction_vectors])
        Y = numpy.concatenate([Y for _ in range(5)], axis=0)
        return X, Y

    from sklearn import datasets, svm, naive_bayes
    from sklearn import cross_validation, preprocessing
    SPOKEN_WORDS = True
    MNIST = False
    DIGITS = False
    NUDGE_DIGITS = True
    FACES = False
    TWENTYNEWSGROUPS = False
    VERBOSE = True
    SCALE = True
    PLOT = True

    def train_models(x_train, y_train, x_test, y_test, n_features, n_outs,
            use_dropout=False, n_epochs=100, numpy_rng=None, # TODO 200+ epochs
            svms=False, nb=False, deepnn=True, name=''):
        if svms:
            print("Linear SVM")
            classifier = svm.SVC(gamma=0.001)
            print(classifier)
            classifier.fit(x_train, y_train)
            print("score: %f" % classifier.score(x_test, y_test))

            print("RBF-kernel SVM")
            classifier = svm.SVC(kernel='rbf', class_weight='auto')
            print(classifier)
            classifier.fit(x_train, y_train)
            print("score: %f" % classifier.score(x_test, y_test))

        if nb:
            print("Multinomial Naive Bayes")
            classifier = naive_bayes.MultinomialNB()
            print(classifier)
            classifier.fit(x_train, y_train)
            print("score: %f" % classifier.score(x_test, y_test))

        if deepnn:
            import warnings
            warnings.filterwarnings("ignore")  # TODO remove

            if use_dropout:
                n_epochs *= 4
                pass

            def new_dnn(dropout=False):
                if dropout:
                    print("Dropout DNN")
                    return DropoutNet(numpy_rng=numpy_rng, n_ins=n_features,
                        #layers_types=[ReLU, ReLU, ReLU, ReLU, LogisticRegression],
                        layers_types=[SoftPlus, SoftPlus, SoftPlus, SoftPlus, LogisticRegression],
                        layers_sizes=[2000, 2000, 2000, 2000],
                        dropout_rates=[0.2, 0.5, 0.5, 0.5, 0.5],
                        n_outs=n_outs,
                        max_norm=4.,
                        fast_drop=False,
                        debugprint=0)
                else:
                    print("Simple (regularized) DNN")
                    return RegularizedNet(numpy_rng=numpy_rng, n_ins=n_features,
                        #layers_types=[LogisticRegression],
                        #layers_sizes=[],
                        #layers_types=[ReLU, ReLU, ReLU, LogisticRegression],
                        #layers_types=[SoftPlus, SoftPlus, SoftPlus, LogisticRegression],
                        #layers_sizes=[1000, 1000, 1000],
                        layers_types=[ReLU, LogisticRegression],
                        layers_sizes=[200],
                        n_outs=n_outs,
                        #L1_reg=0.001/x_train.shape[0],
                        #L2_reg=0.001/x_train.shape[0],
                        L1_reg=0.,
                        L2_reg=1./x_train.shape[0],
                        max_norm=0.,
                        debugprint=0)

            import matplotlib.pyplot as plt
            plt.figure()
            ax1 = plt.subplot(221)
            ax2 = plt.subplot(222)
            ax3 = plt.subplot(223)
            ax4 = plt.subplot(224)  # TODO updates of the weights
            methods = ['adadelta']
            for method in methods:
                dnn = new_dnn(use_dropout)
                print dnn
                dnn.fit(x_train, y_train, max_epochs=n_epochs, method=method, verbose=VERBOSE, plot=PLOT)
                test_error = dnn.score(x_test, y_test)
                print("score: %f" % (1. - test_error))
                ax1.plot(numpy.log10(dnn._costs), label=method)
                #ax2.plot(numpy.log10(dnn._train_errors), label=method)
                #ax3.plot(numpy.log10(dnn._dev_errors), label=method)
                ax2.plot(dnn._train_errors, label=method)
                ax3.plot(dnn._dev_errors, label=method)
                #ax4.plot(dnn._updates, label=method) TODO
                ax4.plot([test_error for _ in range(10)], label=method)
            ax1.set_xlabel('epoch')
            ax1.set_ylabel('cost (log10)')
            ax2.set_xlabel('epoch')
            ax2.set_ylabel('train error')
            ax3.set_xlabel('epoch')
            ax3.set_ylabel('dev error')
            ax4.set_ylabel('test error')
            plt.legend()
            plt.savefig('training_log' + name + '.png')

    if MNIST:
        from sklearn.datasets import fetch_mldata
        mnist = fetch_mldata('MNIST original')
        X = numpy.asarray(mnist.data, dtype='float32')
        if SCALE:
            #X = preprocessing.scale(X)
            X /= 255.
        y = numpy.asarray(mnist.target, dtype='int32')
        #target_names = mnist.target_names
        print("Total dataset size:")
        print("n samples: %d" % X.shape[0])
        print("n features: %d" % X.shape[1])
        print("n classes: %d" % len(set(y)))
        x_train, x_test, y_train, y_test = cross_validation.train_test_split(
                    X, y, test_size=0.2, random_state=42)

        train_models(x_train, y_train, x_test, y_test, X.shape[1],
                     len(set(y)), numpy_rng=numpy.random.RandomState(123),
                     name='MNIST')

    if DIGITS:
        digits = datasets.load_digits()
        data = numpy.asarray(digits.data, dtype='float32')
        target = numpy.asarray(digits.target, dtype='int32')
        x = data
        y = target
        if NUDGE_DIGITS:
            x, y = nudge_dataset(x, y)
        if SCALE:
            x = preprocessing.scale(x)
        x_train, x_test, y_train, y_test = cross_validation.train_test_split(
                x, y, test_size=0.2, random_state=42)
        train_models(x_train, y_train, x_test, y_test, x.shape[1],
                     len(set(target)), numpy_rng=numpy.random.RandomState(123),
                     name='digits')

    if FACES:
        import logging
        logging.basicConfig(level=logging.INFO,
                            format='%(asctime)s %(message)s')
        lfw_people = datasets.fetch_lfw_people(min_faces_per_person=70,
                                               resize=0.4)
        X = numpy.asarray(lfw_people.data, dtype='float32')
        if SCALE:
            X = preprocessing.scale(X)
        y = numpy.asarray(lfw_people.target, dtype='int32')
        target_names = lfw_people.target_names
        print("Total dataset size:")
        print("n samples: %d" % X.shape[0])
        print("n features: %d" % X.shape[1])
        print("n classes: %d" % target_names.shape[0])
        x_train, x_test, y_train, y_test = cross_validation.train_test_split(
                    X, y, test_size=0.2, random_state=42)

        train_models(x_train, y_train, x_test, y_test, X.shape[1],
                     len(set(y)), numpy_rng=numpy.random.RandomState(123),
                     name='faces')

    if TWENTYNEWSGROUPS:
        from sklearn.feature_extraction.text import TfidfVectorizer
        newsgroups_train = datasets.fetch_20newsgroups(subset='train')
        vectorizer = TfidfVectorizer(encoding='latin-1', max_features=10000)
        #vectorizer = HashingVectorizer(encoding='latin-1')
        x_train = vectorizer.fit_transform(newsgroups_train.data)
        x_train = numpy.asarray(x_train.todense(), dtype='float32')
        y_train = numpy.asarray(newsgroups_train.target, dtype='int32')
        newsgroups_test = datasets.fetch_20newsgroups(subset='test')
        x_test = vectorizer.transform(newsgroups_test.data)
        x_test = numpy.asarray(x_test.todense(), dtype='float32')
        y_test = numpy.asarray(newsgroups_test.target, dtype='int32')
        train_models(x_train, y_train, x_test, y_test, x_train.shape[1],
                     len(set(y_train)),
                     numpy_rng=numpy.random.RandomState(123),
                     svms=False, nb=True, deepnn=True,
                     name='20newsgroups')

    if SPOKEN_WORDS:
        # words done by "say", shapes of their filterbanks
        #>>> shapes
        #array([[62, 40],
        #       [65, 40],
        #       [58, 40],
        #       ...,
        #       [85, 40],
        #       [79, 40],
        #       [51, 40]])
        #>>> shapes.mean(axis=0)
        #array([ 70.87751196,  40.        ])
        #>>> shapes.std(axis=0)
        #array([ 12.94580736,   0.        ])
        #>>> shapes.min(axis=0)
        #array([39, 40])
        words_fbanks = numpy.load("all_words_pascal1k.npz")
        n_tokens = len([k for k in words_fbanks.keys()])
        lexicon = set([w.split('_')[1] for w in words_fbanks.keys()])
        lexicon = [w for w in lexicon]  # we need an ordered collection
        n_words = len(lexicon)
        all_fbanks = numpy.concatenate([v for _, v in words_fbanks.iteritems()])
        print all_fbanks.shape
        mean = all_fbanks.mean(axis=0)
        print mean.shape
        std = all_fbanks.std(axis=0)
        print std.shape
        # take 69 fbanks in the middle of the word and pad with 0s if needed
        X = numpy.zeros((n_tokens, 40*STACKSIZE), dtype='float32')
        y = numpy.zeros(n_tokens, dtype='int32')
        for i, (swf, fb) in enumerate(words_fbanks.iteritems()):
            spkr, word, _ = swf.split('_')
            l = fb.shape[0]
            m = l/2
            s = max(0, m - ((STACKSIZE-1) / 2))
            e = min(l-1, m + ((STACKSIZE-1) / 2))
            tmp = (fb - mean) / std
            tmp = tmp[s:e+1].flatten()
            diff = 40*STACKSIZE - tmp.shape[0]
            if not diff:
                X[i] = tmp
            else:
                X[i][diff/2:-diff/2] = tmp
            y[i] = lexicon.index(word)
        # train the DNN, with the training set as test set if let in this form:
        train_models(X, y, X, y, X.shape[1],
                     len(set(y)),
                     numpy_rng=numpy.random.RandomState(123),
                     svms=False, nb=False, deepnn=True,
                     name='spoken_words')


# Copyright 2016 Google Inc. 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.

"""Implements usefule CSV utilities."""
from __future__ import absolute_import
from __future__ import unicode_literals
from builtins import next
from builtins import str as newstr
from builtins import range
from builtins import object

import csv
import os
import pandas as pd
import random
import sys

try:
    from StringIO import StringIO
except ImportError:
    from io import StringIO
import tempfile
import datalab.storage
import datalab.utils

_MAX_CSV_BYTES = 10000000

class Csv(object):
  """Represents a CSV file in GCS or locally with same schema.
  """
  def __init__(self, path, delimiter=b','):
    """Initializes an instance of a Csv instance.
    Args:
      path: path of the Csv file.
      delimiter: the separator used to parse a Csv line.
    """
    self._path = path
    self._delimiter = delimiter

  @property
  def path(self):
    return self._path

  @staticmethod
  def _read_gcs_lines(path, max_lines=None):
    return datalab.storage.Item.from_url(path).read_lines(max_lines)

  @staticmethod
  def _read_local_lines(path, max_lines=None):
    lines = []
    for line in open(path):
      if max_lines is not None and len(lines) >= max_lines:
        break
      lines.append(line)
    return lines

  def _is_probably_categorical(self, column):
    if newstr(column.dtype) != 'object':
      # only string types (represented in DataFrame as object) can potentially be categorical
      return False
    if len(max(column, key=lambda p: len(newstr(p)))) > 100:
      return False  # value too long to be a category
    if len(set(column)) > 100:
      return False  # too many unique values to be a category
    return True

  def browse(self, max_lines=None, headers=None):
    """Try reading specified number of lines from the CSV object.
    Args:
      max_lines: max number of lines to read. If None, the whole file is read
      headers: a list of strings as column names. If None, it will use "col0, col1..."
    Returns:
      A pandas DataFrame with the schema inferred from the data.
    Raises:
      Exception if the csv object cannot be read or not enough lines to read, or the
      headers size does not match columns size.
    """
    if self.path.startswith('gs://'):
      lines = Csv._read_gcs_lines(self.path, max_lines)
    else:
      lines = Csv._read_local_lines(self.path, max_lines)
    if len(lines) == 0:
      return pd.DataFrame(columns=headers)
    columns_size = len(next(csv.reader([lines[0]], delimiter=self._delimiter)))
    if headers is None:
      headers = ['col' + newstr(e) for e in range(columns_size)]
    if len(headers) != columns_size:
      raise Exception('Number of columns in CSV do not match number of headers')
    buf = StringIO()
    for line in lines:
      buf.write(line)
      buf.write('\n')
    buf.seek(0)
    df = pd.read_csv(buf, names=headers, delimiter=self._delimiter)
    for key, col in df.iteritems():
      if self._is_probably_categorical(col):
        df[key] = df[key].astype('category')
    return df

  def _create_federated_table(self, skip_header_rows):
    import datalab.bigquery as bq
    df = self.browse(1, None)
    # read each column as STRING because we only want to sample rows.
    schema_train = bq.Schema([{'name': name, 'type': 'STRING'} for name in df.keys()])
    options = bq.CSVOptions(skip_leading_rows=(1 if skip_header_rows is True else 0))
    return bq.FederatedTable.from_storage(self.path,
                                          csv_options=options,
                                          schema=schema_train,
                                          max_bad_records=0)

  def _get_gcs_csv_row_count(self, federated_table):
    import datalab.bigquery as bq
    results = bq.Query('SELECT count(*) from data',
                       data_sources={'data': federated_table}).results()
    return results[0].values()[0]

  def sample_to(self, count, skip_header_rows, strategy, target):
    """Sample rows from GCS or local file and save results to target file.

    Args:
      count: number of rows to sample. If strategy is "BIGQUERY", it is used as approximate number.
      skip_header_rows: whether to skip first row when reading from source.
      strategy: can be "LOCAL" or "BIGQUERY". If local, the sampling happens in local memory,
          and number of resulting rows matches count. If BigQuery, sampling is done
          with BigQuery in cloud, and the number of resulting rows will be approximated to
          count.
      target: The target file path, can be GCS or local path.
    Raises:
      Exception if strategy is "BIGQUERY" but source is not a GCS path.
    """
    # TODO(qimingj) Add unit test
    # Read data from source into DataFrame.

    if sys.version_info.major > 2:
      xrange = range  # for python 3 compatibility

    if strategy == 'BIGQUERY':
      import datalab.bigquery as bq
      if not self.path.startswith('gs://'):
        raise Exception('Cannot use BIGQUERY if data is not in GCS')
      federated_table = self._create_federated_table(skip_header_rows)
      row_count = self._get_gcs_csv_row_count(federated_table)
      query = bq.Query('SELECT * from data', data_sources={'data': federated_table})
      sampling = bq.Sampling.random(count * 100 / float(row_count))
      sample = query.sample(sampling=sampling)
      df = sample.to_dataframe()
    elif strategy == 'LOCAL':
      local_file = self.path
      if self.path.startswith('gs://'):
        local_file = tempfile.mktemp()
        datalab.utils.gcs_copy_file(self.path, local_file)
      with open(local_file) as f:
        row_count = sum(1 for line in f)
      start_row = 1 if skip_header_rows is True else 0
      skip_count = row_count - count - 1 if skip_header_rows is True else row_count - count
      skip = sorted(random.sample(xrange(start_row, row_count), skip_count))
      header_row = 0 if skip_header_rows is True else None
      df = pd.read_csv(local_file, skiprows=skip, header=header_row, delimiter=self._delimiter)
      if self.path.startswith('gs://'):
        os.remove(local_file)
    else:
      raise Exception('strategy must be BIGQUERY or LOCAL')
    # Write to target.
    if target.startswith('gs://'):
      with tempfile.NamedTemporaryFile() as f:
        df.to_csv(f, header=False, index=False)
        f.flush()
        datalab.utils.gcs_copy_file(f.name, target)
    else:
      with open(target, 'w') as f:
        df.to_csv(f, header=False, index=False, sep=str(self._delimiter))


import os

import dill
import numpy as np
from sklearn.model_selection import KFold

from marseille.custom_logging import logging
from marseille.datasets import get_dataset_loader, load_embeds
from marseille.io import cache_fname
from marseille.argrnn import ArgumentLSTM

def argrnn_cv_score(dataset, dynet_weight_decay, mlp_dropout,
                    rnn_dropout, prop_layers, class_weight, constraints,
                    compat_features, second_order):

    fn = cache_fname("argrnn_cv_score", (dataset, dynet_weight_decay,
                                         mlp_dropout, rnn_dropout, prop_layers,
                                         class_weight, constraints,
                                         compat_features, second_order))

    if os.path.exists(fn):
        logging.info("Cached file already exists.")
        with open(fn, "rb") as f:
            return dill.load(f)

    load, ids = get_dataset_loader(dataset, split="train")
    embeds = load_embeds(dataset)

    grandparent_layers = 1 if second_order and dataset == 'ukp' else 0
    coparent_layers = 1 if second_order else 0
    sibling_layers = 1 if second_order and dataset == 'cdcp' else 0

    scores = []
    all_Y_pred = []
    score_at_iter = [10, 25, 50, 75, 100]

    n_folds = 5 if dataset == 'ukp' else 3

    for k, (tr, val) in enumerate(KFold(n_folds).split(ids)):
        docs_train = list(load(ids[tr]))
        docs_val = list(load(ids[val]))

        Y_train = [doc.label for doc in docs_train]
        Y_val = [doc.label for doc in docs_val]

        rnn = ArgumentLSTM(lstm_dropout=rnn_dropout,
                           mlp_dropout=mlp_dropout,
                           compat_features=compat_features,
                           constraints=constraints,
                           prop_mlp_layers=prop_layers,
                           coparent_layers=coparent_layers,
                           grandparent_layers=grandparent_layers,
                           sibling_layers=sibling_layers,
                           class_weight=class_weight,
                           second_order_multilinear=True,
                           max_iter=100,
                           score_at_iter=score_at_iter,
                           n_mlp=128,
                           n_lstm=128,
                           lstm_layers=2,
                           link_mlp_layers=1,
                           embeds=embeds,
                           exact_inference=False,
                           link_bilinear=True)

        rnn.fit(docs_train, Y_train, docs_val, Y_val)
        Y_val_pred = rnn.predict(docs_val)
        all_Y_pred.extend(Y_val_pred)
        scores.append(rnn.scores_)

    with open(fn, "wb") as f:
        dill.dump((scores, score_at_iter, all_Y_pred), f)

    return scores, score_at_iter, all_Y_pred

if __name__ == '__main__':
    from docopt import docopt

    usage = """
    Usage:
        exp_rnn (cdcp|ukp) [\
--dynet-seed N --dynet-weight-decay N --dynet-mem N --prop-layers=N \
--rnn-dropout=N --mlp-dropout=N --balanced --constraints --strict \
--compat-features --second-order]

    Options:
        --dynet-seed=N          random number generator seed for dynet library
        --dynet-weight-decay=N  global weight decay amount for dynet library
        --dynet-mem=N           memory pool size for dynet
        --prop-layers=N         number of prop classifier layers. [default: 2]
        --rnn-dropout=N         dropout ratio in lstm. [default: 0.0]
        --mlp-dropout=N         dropout ratio in mlp. [default: 0.1]
        --balanced              whether to reweight class costs by freq
        --constraints           whether to constrain the decoding
        --strict                whether to use strict domain constraints
        --compat-features       whether to use features for compat factors
        --second-order          whether to use coparent / grandpa / siblings
    """

    args = docopt(usage)

    dataset = 'cdcp' if args['cdcp'] else 'ukp'
    prop_layers = int(args['--prop-layers'])
    rnn_dropout = float(args['--rnn-dropout'])
    mlp_dropout = float(args['--mlp-dropout'])
    cw = 'balanced' if args['--balanced'] else None

    if args['--constraints']:
        constraints = dataset
        if args['--strict']:
            constraints += '+strict'
    else:
        constraints = ""

    scores, score_at_iter, _ = argrnn_cv_score(dataset,
                                               args['--dynet-weight-decay'],
                                               mlp_dropout,
                                               rnn_dropout,
                                               prop_layers,
                                               cw,
                                               constraints,
                                               args['--compat-features'],
                                               args['--second-order'])

    for iter, score in zip(score_at_iter, np.mean(scores, axis=0)):
        print("iter={} "
              "Link: {:.3f}/{:.3f} "
              "Node: {:.3f}/{:.3f} "
              "accuracy {:.3f}".format(iter, *score),
        )


import unittest
import nose

from numpy.testing import assert_almost_equal, assert_allclose
from numpy.testing.decorators import slow
from pandas.util.testing import (assert_series_equal, assert_frame_equal)

import os
from tempfile import NamedTemporaryFile
import pandas as pd
from pandas import DataFrame, Series

import mr
import sqlite3

path, _ = os.path.split(os.path.abspath(__file__))

class TestFeatureSaving(unittest.TestCase):

    def setUp(self):
        self.db_conn = sqlite3.connect(':memory:')
        directory = os.path.join(path, 'video', 'image_sequence')
        self.v = mr.ImageSequence(directory)
        self.PARAMS = (11, 3000)
        with NamedTemporaryFile() as temp:
            self.expected = mr.batch(self.v[[0, 1]], *self.PARAMS,
                                     meta=temp.name)

    def test_sqlite(self):
        with NamedTemporaryFile() as temp:
            f = mr.batch(self.v[[0, 1]], *self.PARAMS, conn=self.db_conn,
                     sql_flavor='sqlite', table='features', meta=temp.name)
        assert_frame_equal(f, self.expected)

    def test_HDFStore(self):
        STORE_NAME = 'temp_for_testing.h5'
        if os.path.isfile(STORE_NAME):
            os.remove(STORE_NAME)
        try:
            store = pd.HDFStore(STORE_NAME)
        except:
            nose.SkipTest('Cannot make an HDF5 file. Skipping')
        else:
            with NamedTemporaryFile() as temp:
                f = mr.batch(self.v[[0, 1]], *self.PARAMS, store=store,
                             table='features', meta=temp.name)
            assert_frame_equal(f.reset_index(drop=True),
                           self.expected.reset_index(drop=True))
            os.remove(STORE_NAME)


# Copyright 2014-2015 Yandex LLC and contributors <https://yandex.com/>
#
# 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.

from __future__ import division, print_function, absolute_import
from rep.test.test_estimators import check_classifier, check_regression, check_params, \
    generate_classification_data, check_classification_reproducibility
from rep.estimators.pybrain import PyBrainClassifier, PyBrainRegressor
from sklearn.ensemble import BaggingClassifier
from rep.estimators import SklearnClassifier

__author__ = 'Artem Zhirokhov'

classifier_params = {
    'has_staged_pp': False,
    'has_importances': False,
    'supports_weight': False
}

regressor_params = {
    'has_staged_predictions': False,
    'has_importances': False,
    'supports_weight': False
}

def test_pybrain_params():
    check_params(PyBrainClassifier, layers=[1, 2], epochs=5, use_rprop=True, hiddenclass=['LinearLayer'])
    check_params(PyBrainRegressor, layers=[1, 2], epochs=5, etaplus=1.3, hiddenclass=['LinearLayer'], learningrate=0.1)

def test_pybrain_classification():
    clf = PyBrainClassifier(epochs=2)
    check_classifier(clf, **classifier_params)
    check_classifier(PyBrainClassifier(epochs=-1, continue_epochs=1, layers=[]), **classifier_params)
    check_classifier(PyBrainClassifier(epochs=2, layers=[5, 2]), **classifier_params)

def test_pybrain_reproducibility():
    try:
        import numpy
        X, y, _ = generate_classification_data()
        clf1 = PyBrainClassifier(layers=[4], epochs=2).fit(X, y)
        clf2 = PyBrainClassifier(layers=[4], epochs=2).fit(X, y)
        print(clf1.predict_proba(X)-clf2.predict_proba(X))
        assert numpy.allclose(clf1.predict_proba(X), clf2.predict_proba(X)), 'different predicitons'
        check_classification_reproducibility(clf1, X, y)
    except:
        # This test fails. Because PyBrain can't reproduce training.
        pass

def test_pybrain_Linear_MDLSTM():
    check_classifier(PyBrainClassifier(epochs=2, layers=[10, 2], hiddenclass=['LinearLayer', 'MDLSTMLayer']),
                     **classifier_params)
    check_regression(PyBrainRegressor(epochs=3, layers=[10, 2], hiddenclass=['LinearLayer', 'MDLSTMLayer']),
                     **regressor_params)

def test_pybrain_SoftMax_Tanh():
    check_classifier(PyBrainClassifier(epochs=2, layers=[10, 5, 2], hiddenclass=['SoftmaxLayer', 'SoftmaxLayer', 'TanhLayer'], use_rprop=True),
                     **classifier_params)
    check_regression(PyBrainRegressor(epochs=2, layers=[10, 5, 2], hiddenclass=['SoftmaxLayer', 'TanhLayer', 'TanhLayer']),
                     **regressor_params)

def pybrain_test_partial_fit():
    clf = PyBrainClassifier(layers=[4], epochs=2)
    X, y, _ = generate_classification_data()
    clf.partial_fit(X, y)
    clf.partial_fit(X[:2], y[:2])

def test_pybrain_multi_classification():
    check_classifier(PyBrainClassifier(), n_classes=4, **classifier_params)

def test_pybrain_regression():
    check_regression(PyBrainRegressor(), **regressor_params)

def test_pybrain_multi_regression():
    check_regression(PyBrainRegressor(), n_targets=4, **regressor_params)

def test_simple_stacking_pybrain():
    base_pybrain = PyBrainClassifier()
    base_bagging = BaggingClassifier(base_estimator=base_pybrain, n_estimators=3)
    check_classifier(SklearnClassifier(clf=base_bagging), **classifier_params)


import numpy as np
import pandas as pd

from statsmodels import regression
import statsmodels.api as sm
import matplotlib.pyplot as plt
import math

import pandas_datareader.data as web
from datetime import datetime

def Starter_Lec7():
    start = datetime(2014, 1, 1)
    end = datetime(2015, 1, 1)

    asset = web.DataReader("TSLA","yahoo",start,end)
    asset_closingPrice = asset['Close']

    benchmark = web.DataReader("SPY","yahoo",start,end)
    benchmark_closingPrice = benchmark['Close']

    r_a = asset_closingPrice.pct_change()[1:]
    r_b = benchmark_closingPrice.pct_change()[1:]

    modelSummary = linreg(r_a,r_b)
    print("{0} {1} \n\n".format(modelSummary,type(modelSummary)))

def linreg(X,Y):
    #running linear regression
    X = sm.add_constant(X)
    model = regression.linear_model.OLS(Y,X).fit()

    a = model.params[0]
    b = model.params[1]

    X = pd.DataFrame(X, columns=['Close'])     #Y_CMT Neat trick to extract columns from a pandas dataframe

    # Return summary of the regression and plot results
    X2 = np.linspace(float(X.min()), float(X.max()), 100)
    Y_hat = X2 * b + a

    plt.scatter(X, Y, alpha=0.3) # Plot the raw data

    plt.plot(X2, Y_hat, 'r', alpha=0.9)  # Add the regression line, colored in red
    plt.xlabel('X Value')
    plt.ylabel('Y Value')
    plt.show()
    return model.summary()

def TestPlotting():
    N = 8
    y = np.zeros(N)

    x1 = np.linspace(0, 10, N, endpoint=True)
    x2 = np.linspace(0, 10, N, endpoint=False)

    plt.plot(x1, y, 'o')
    plt.plot(x2, y + 0.5, 'o')
    plt.ylim([-0.5, 1])
    plt.show()

def NumpyMatrix():
    array1 = np.matrix([[1,2,3],[4,5,6],[7,8,9]])
    print("{0} {1} \n\n".format(array1[:,2],type(array1)))

    array1 = array1[:,2]
    print("{0} {1} \n\n".format(array1,type(array1)))


"""
Copyright (C) 2013 Travis DeWolf

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program 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 General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

from pydmps.dmp import DMPs

import numpy as np

class DMPs_rhythmic(DMPs):
    """An implementation of discrete DMPs"""

    def __init__(self, **kwargs):
        """
        """

        # call super class constructor
        super(DMPs_rhythmic, self).__init__(pattern="rhythmic", **kwargs)

        self.gen_centers()

        # set variance of Gaussian basis functions
        # trial and error to find this spacing
        self.h = np.ones(self.n_bfs) * self.n_bfs  # 1.75

        self.check_offset()

    def gen_centers(self):
        """Set the centre of the Gaussian basis
        functions be spaced evenly throughout run time"""

        c = np.linspace(0, 2 * np.pi, self.n_bfs + 1)
        c = c[0:-1]
        self.c = c

    def gen_front_term(self, x, dmp_num):
        """Generates the front term on the forcing term.
        For rhythmic DMPs it's non-diminishing, so this
        function is just a placeholder to return 1.

        x float: the current value of the canonical system
        dmp_num int: the index of the current dmp
        """

        if isinstance(x, np.ndarray):
            return np.ones(x.shape)
        return 1

    def gen_goal(self, y_des):
        """Generate the goal for path imitation.
        For rhythmic DMPs the goal is the average of the
        desired trajectory.

        y_des np.array: the desired trajectory to follow
        """

        goal = np.zeros(self.n_dmps)
        for n in range(self.n_dmps):
            num_idx = ~np.isnan(y_des[n])  # ignore nan's when calculating goal
            goal[n] = 0.5 * (y_des[n, num_idx].min() + y_des[n, num_idx].max())

        return goal

    def gen_psi(self, x):
        """Generates the activity of the basis functions for a given
        canonical system state or path.

        x float, array: the canonical system state or path
        """

        if isinstance(x, np.ndarray):
            x = x[:, None]
        return np.exp(self.h * (np.cos(x - self.c) - 1))

    def gen_weights(self, f_target):
        """Generate a set of weights over the basis functions such
        that the target forcing term trajectory is matched.

        f_target np.array: the desired forcing term trajectory
        """

        # calculate x and psi
        x_track = self.cs.rollout()
        psi_track = self.gen_psi(x_track)

        # efficiently calculate BF weights using weighted linear regression
        for d in range(self.n_dmps):
            for b in range(self.n_bfs):
                self.w[d, b] = np.dot(psi_track[:, b], f_target[:, d]) / (
                    np.sum(psi_track[:, b]) + 1e-10
                )

# ==============================
# Test code
# ==============================
if __name__ == "__main__":
    import matplotlib.pyplot as plt

    # test normal run
    dmp = DMPs_rhythmic(n_dmps=1, n_bfs=10, w=np.zeros((1, 10)))
    y_track, dy_track, ddy_track = dmp.rollout()

    plt.figure(1, figsize=(6, 3))
    plt.plot(np.ones(len(y_track)) * dmp.goal, "r--", lw=2)
    plt.plot(y_track, lw=2)
    plt.title("DMP system - no forcing term")
    plt.xlabel("time (ms)")
    plt.ylabel("system trajectory")
    plt.legend(["goal", "system state"], loc="lower right")
    plt.tight_layout()

    # test imitation of path run
    plt.figure(2, figsize=(6, 4))
    n_bfs = [10, 30, 50, 100, 10000]

    # a straight line to target
    path1 = np.sin(np.arange(0, 2 * np.pi, 0.01) * 5)
    # a strange path to target
    path2 = np.zeros(path1.shape)
    path2[int(len(path2) / 2.0) :] = 0.5

    for ii, bfs in enumerate(n_bfs):
        dmp = DMPs_rhythmic(n_dmps=2, n_bfs=bfs)

        dmp.imitate_path(y_des=np.array([path1, path2]))
        y_track, dy_track, ddy_track = dmp.rollout()

        plt.figure(2)
        plt.subplot(211)
        plt.plot(y_track[:, 0], lw=2)
        plt.subplot(212)
        plt.plot(y_track[:, 1], lw=2)

    plt.subplot(211)
    a = plt.plot(path1, "r--", lw=2)
    plt.title("DMP imitate path")
    plt.xlabel("time (ms)")
    plt.ylabel("system trajectory")
    plt.legend([a[0]], ["desired path"], loc="lower right")
    plt.subplot(212)
    b = plt.plot(path2, "r--", lw=2)
    plt.title("DMP imitate path")
    plt.xlabel("time (ms)")
    plt.ylabel("system trajectory")
    plt.legend(["%i BFs" % i for i in n_bfs], loc="lower right")

    plt.tight_layout()
    plt.show()