"""
   Copyright 2019 Riley John Murray

   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.
"""
import numpy as np


def mat_times_vecvar(mat, vecvar):
    """

    :param mat: a numpy ndarray of shape (m, n).
    :param vecvar: a coniclifts Variable of shape (n,) or (n, 1).

    :return: A_vals, A_rows, A_cols so that the coniclifts Expression
        con = (mat @ vecvar >= 0) would compile to
        A_vals, A_rows, A_cols, np.zeros((m,)), [].
    """
    ids = vecvar.scalar_variable_ids
    A_rows = np.tile(np.arange(mat.shape[0]), reps=mat.shape[1])
    A_cols = np.repeat(ids, mat.shape[0]).tolist()
    A_vals = mat.ravel(order='F').tolist()  # stack columns, then tolist
    return A_vals, A_rows, A_cols


def mat_times_vecvar_plus_mat_times_vecvar(mat1, vecvar1, mat2, vecvar2):
    """

    :param mat1: a numpy ndarray of shape (m, n1).
    :param vecvar1: a coniclifts Variable of shape (n1,) or (n1, 1).
    :param mat2: a numpy ndarray of shape (m, n2)
    :param vecvar2: a coniclifts Variable of shape (n2,) or (n2, 1).

    :return: A_vals, A_rows, A_cols so that the coniclifts Expression
        expr = mat1 @ vecvar1 + mat2 @ vecvar2 would have "expr >= 0"
        compile to A_vals, A_rows, A_cols, np.zeros((m,)), [].
    """
    mat = np.hstack([mat1, mat2])
    ids = [v.id for v in vecvar1.scalar_variables()]
    ids += [v.id for v in vecvar2.scalar_variables()]
    A_rows = np.tile(np.arange(mat.shape[0]), reps=mat.shape[1])
    A_cols = np.repeat(ids, mat.shape[0]).tolist()
    A_vals = mat.ravel(order='F').tolist()  # stack columns, then tolist
    return A_vals, A_rows, A_cols


def mat_times_vecvar_minus_vecvar(mat, vecvar1, vecvar2):
    """

    :param mat: a numpy ndarray of shape (m, n).
    :param vecvar1: a coniclifts Variable of shape (n,) or (n, 1).
    :param vecvar2: a coniclifts Variable of shape (m,) or (m, 1).

    :return: A_vals, A_rows, A_cols so that the coniclifts Expression
        expr = mat @ vecvar1 - vecvar2 would have "expr >= 0"
        compile to A_vals, A_rows, A_cols, np.zeros((m,)), [].
    """
    A_vals, A_rows, A_cols = [], [], []
    num_rows = mat.shape[0]
    vec1_ids = vecvar1.scalar_variable_ids
    vec2_ids = vecvar2.scalar_variable_ids
    for j in range(num_rows):
        A_vals += mat[j, :].tolist()
        A_vals.append(-1)
        A_cols += vec1_ids
        A_cols.append(vec2_ids[j])
        A_rows += [j] * (mat.shape[1] + 1)
    A_rows = np.array(A_rows)
    return A_vals, A_rows, A_cols


def mat_times_vecvar_plus_vec_times_singlevar(mat, vecvar, vec, singlevar):
    """

    :param mat: a numpy ndarray of shape (m, n).
    :param vecvar: a coniclifts Variable of shape (n,) or (n, 1)
    :param vec: a numpy ndarray of shape (m,) or (m, 1).
    :param singlevar: a coniclifts Variable of size 1.
    :return:

    Return A_rows, A_cols, A_vals as if they were generated by a compile() call on
    con = (mat @ vecvar + vec * singlevar >= 0).
    """
    var_ids = vecvar.scalar_variable_ids.copy() + [singlevar.scalar_variables()[0].id]
    ids = np.array(var_ids)
    mat = np.hstack([mat, np.reshape(vec, (-1, 1))])
    A_rows = np.tile(np.arange(mat.shape[0]), reps=mat.shape[1])
    A_cols = np.repeat(ids, mat.shape[0]).tolist()
    A_vals = mat.ravel(order='F').tolist()  # stack columns, then tolist
    return A_vals, A_rows, A_cols


def columns_sum_to_vec(mat, vec):
    A_rows, A_cols, A_vals = [], [], []
    m = mat.shape[0]
    if m != vec.size:
        raise RuntimeError('Incompatible dimensions.')
    b = np.zeros(m,)
    for i in range(m):
        # update cols and data to reflect addition of elements in ith row of mat
        svs = mat[i, :].scalar_variables()
        A_cols += [sv.id for sv in svs]
        A_vals += [1] * len(svs)
        # update cols and data to reflect addition of elements from ith element of vec
        #   ith element of vec is a ScalarExpression!
        id2co = [(a.id, co) for a, co in vec[i].atoms_to_coeffs.items()]
        A_cols += [aid for aid, _ in id2co]
        A_vals += [-co for _, co in id2co]  # we are subtracting, after all.
        # update rows with appropriate number of "i"s.
        A_rows += [i] * (len(svs) + len(id2co))
        # update b
        b[i] -= vec[i].offset
    A_rows = np.array(A_rows)
    return A_vals, A_rows, A_cols, b


