#!/usr/bin/env python
"""
Summary
-------
Useful matrix/modulus operations for mrg32k3a generator.

Listing
-------
mat33_mat31_mult
mat33_mat33_mult
mat31_mod
mat33_mod
mat33_mat33_mod
mat33_power_mod
"""


def mat33_mat31_mult(A, b):
    """
    Multiply a 3x3 matrix with a 3x1 matrix.

    Arguments
    ---------
    A : list of list of float
        3x3 matrix
    b : list of float
        3x1 matrix

    Returns
    -------
    res : list of float
        3x1 matrix
    """
    res = [0, 0, 0]
    r3 = range(3)
    for i in r3:
        res[i] = sum([A[i][j] * b[j] for j in r3])
    return res


def mat33_mat33_mult(A, B):
    """
    Multiply a 3x3 matrix with a 3x3 matrix.

    Arguments
    ---------
    A : list of list of float
        3x3 matrix
    B : list of list of float
        3x3 matrix

    Returns
    -------
    res : list of float
        3x3 matrix
    """
    res = [[0, 0, 0],
           [0, 0, 0],
           [0, 0, 0]
           ]
    r3 = range(3)
    for i in r3:
        for j in r3:
            res[i][j] = sum([A[i][k] * B[k][j] for k in r3])
    return res


def mat31_mod(b, m):
    """
    Compute moduli of a 3x1 matrix.

    Arguments
    ---------
    b : list of float
        3x1 matrix
    m : float
        modulus

    Returns
    -------
    res : list of float
        3x1 matrix
    """
    res = [0, 0, 0]
    for i in range(3):
        res[i] = int(b[i] - int(b[i] / m) * m)
        # if negative, add back modulus m
        if res[i] < 0:
            res[i] += m
    return res


def mat33_mod(A, m):
    """
    Compute moduli of a 3x3 matrix.

    Arguments
    ---------
    A : list of float
        3x3 matrix
    m : float
        modulus

    Returns
    -------
    res : list of float
        3x3 matrix
    """
    res = [[0, 0, 0],
           [0, 0, 0],
           [0, 0, 0]
           ]
    r3 = range(3)
    for i in r3:
        for j in r3:
            res[i][j] = int(A[i][j] - int(A[i][j] / m) * m)
            # if negative, add back modulus m
            if res[i][j] < 0:
                res[i][j] += m
    return res


def mat33_mat33_mod(A, B, m):
    """
    Compute moduli of a 3x3 matrix x 3x3 matrix product.

    Arguments
    ---------
    A : list of list of float
        3x3 matrix
    B : list of list of float
        3x3 matrix
    m : float
        modulus

    Returns
    -------
    res : list of list of float
        3x3 matrix
    """
    C = mat33_mat33_mult(A, B)
    res = mat33_mod(C, m)
    return res


def mat33_power_mod(A, j, m):
    """
    Compute moduli of a 3x3 matrix power.
    Use divide-and-conquer algorithm described in L'Ecuyer (1990).

    Arguments
    ---------
    A : list of list of float
        3x3 matrix
    j : int
        exponent
    m : float
        modulus

    Returns
    -------
    res : list of list of float
        3x3 matrix
    """
    B = [[1, 0, 0],
         [0, 1, 0],
         [0, 0, 1]
         ]
    while j > 0:
        if (j % 2 == 1):
            B = mat33_mat33_mod(A, B, m)
        A = mat33_mat33_mod(A, A, m)
        j = int(j / 2)
    res = B
    return res
