# cython: boundscheck=False
# Copyright ExplsionAI GmbH, released under BSD.
cimport numpy as np
from . cimport cy
from .cy cimport reals1d_ft, reals2d_ft, float1d_t, float2d_t
from .cy cimport const_reals1d_ft, const_reals2d_ft, const_float1d_t, const_float2d_t
from .cy cimport const_double1d_t, const_double2d_t

import numpy


def axpy(const_reals1d_ft A, double scale=1., np.ndarray out=None):
    if const_reals1d_ft is const_float1d_t:
        if out is None:
            out = numpy.zeros((A.shape[0],), dtype='f')
        B = <float*>out.data
        return out
    elif const_reals1d_ft is const_double1d_t:
        if out is None:
            out = numpy.zeros((A.shape[0],), dtype='d')
        B = <double*>out.data
        with nogil:
            cy.axpyv(cy.NO_CONJUGATE, A.shape[0], scale, &A[0], 1, B, 1)
        return out
    else:
        B = NULL
        raise TypeError("Unhandled fused type")


def batch_axpy(reals2d_ft A, reals1d_ft B, np.ndarray out=None):
    pass


def ger(const_reals2d_ft A, const_reals1d_ft B, double scale=1., np.ndarray out=None):
    if const_reals2d_ft is const_float2d_t and const_reals1d_ft is const_float1d_t:
        if out is None:
            out = numpy.zeros((A.shape[0], B.shape[0]), dtype='f')
        with nogil:
            cy.ger(
                cy.NO_CONJUGATE, cy.NO_CONJUGATE,
                A.shape[0], B.shape[0],
                scale,
                &A[0,0], 1,
                &B[0], 1,
                <float*>out.data, out.shape[1], 1)
        return out
    elif const_reals2d_ft is const_double2d_t and const_reals1d_ft is const_double1d_t:
        if out is None:
            out = numpy.zeros((A.shape[0], B.shape[0]), dtype='d')
        with nogil:
            cy.ger(
                cy.NO_CONJUGATE, cy.NO_CONJUGATE,
                A.shape[0], B.shape[0],
                scale,
                &A[0,0], 1,
                &B[0], 1,
                <double*>out.data, out.shape[1], 1)
        return out
    else:
        C = NULL
        raise TypeError("Unhandled fused type")


def gemm(const_reals2d_ft A, const_reals2d_ft B,
         np.ndarray out=None, bint trans1=False, bint trans2=False,
         double alpha=1., double beta=1.):
    cdef cy.dim_t nM = A.shape[0] if not trans1 else A.shape[1]
    cdef cy.dim_t nK = A.shape[1] if not trans1 else A.shape[0]
    cdef cy.dim_t nN = B.shape[1] if not trans2 else B.shape[0]
    if const_reals2d_ft is const_float2d_t:
        if out is None:
            out = numpy.zeros((nM, nN), dtype='f')
        C = <float*>out.data
        with nogil:
            cy.gemm(
                cy.TRANSPOSE if trans1 else cy.NO_TRANSPOSE,
                cy.TRANSPOSE if trans2 else cy.NO_TRANSPOSE,
                nM, nN, nK,
                alpha,
                &A[0,0], A.shape[1], 1,
                &B[0,0], B.shape[1], 1,
                beta,
                C, out.shape[1], 1)
        return out
    elif const_reals2d_ft is const_double2d_t:
        if out is None:
            out = numpy.zeros((A.shape[0], B.shape[1]), dtype='d')
        C = <double*>out.data
        with nogil:
            cy.gemm(
                cy.TRANSPOSE if trans1 else cy.NO_TRANSPOSE,
                cy.TRANSPOSE if trans2 else cy.NO_TRANSPOSE,
                A.shape[0], B.shape[1], A.shape[1],
                alpha,
                &A[0,0], A.shape[1], 1,
                &B[0,0], B.shape[1], 1,
                beta,
                C, out.shape[1], 1)
        return out
    else:
        C = NULL
        raise TypeError("Unhandled fused type")

def gemv(const_reals2d_ft A, const_reals1d_ft B,
        bint trans1=False, double alpha=1., double beta=1.,
        np.ndarray out=None):
    if const_reals1d_ft is const_float1d_t and const_reals2d_ft is const_float2d_t:
        if out is None:
            out = numpy.zeros((A.shape[0],), dtype='f')
        with nogil:
            cy.gemv(
                cy.TRANSPOSE if trans1 else cy.NO_TRANSPOSE,
                cy.NO_CONJUGATE,
                A.shape[0], A.shape[1],
                alpha,
                &A[0,0], A.shape[1], 1,
                &B[0], 1,
                beta,
                <float*>out.data, 1)
        return out
    elif const_reals1d_ft is const_double1d_t and const_reals2d_ft is const_double2d_t:
        if out is None:
            out = numpy.zeros((A.shape[0],), dtype='d')
        with nogil:
            cy.gemv(
                cy.TRANSPOSE if trans1 else cy.NO_TRANSPOSE,
                cy.NO_CONJUGATE,
                A.shape[0], A.shape[1],
                alpha,
                &A[0,0], A.shape[1], 1,
                &B[0], 1,
                beta,
                <double*>out.data, 1)
        return out
    else:
        raise TypeError("Unhandled fused type")


def dotv(const_reals1d_ft X, const_reals1d_ft Y, bint conjX=False, bint conjY=False):
    if X.shape[0] != Y.shape[0]:
        msg = "Shape mismatch for blis.dotv: (%d,), (%d,)"
        raise ValueError(msg % (X.shape[0], Y.shape[0]))
    return cy.dotv(
        cy.CONJUGATE if conjX else cy.NO_CONJUGATE,
        cy.CONJUGATE if conjY else cy.NO_CONJUGATE,
        X.shape[0], &X[0], &Y[0], 1, 1
    )


def einsum(todo, A, B, out=None):
    if todo == 'a,a->a':
        return axpy(A, B, out=out)
    elif todo == 'a,b->ab':
        return ger(A, B, out=out)
    elif todo == 'a,b->ba':
        return ger(B, A, out=out)
    elif todo == 'ab,a->ab':
        return batch_axpy(A, B, out=out)
    elif todo == 'ab,a->ba':
        return batch_axpy(A, B, trans1=True, out=out)
    elif todo == 'ab,b->a':
        return gemv(A, B, out=out)
    elif todo == 'ab,a->b':
        return gemv(A, B, trans1=True, out=out)
    # The rule here is, look at the first dimension of the output. That must
    # occur in arg1. Set trans1 if it's dimension 2.
    # E.g. bc is output, b occurs in ab, so that must be arg1. So we need
    # trans1=True, to make ba,ac->bc
    elif todo == 'ab,ac->bc':
        return gemm(A, B, trans1=True, trans2=False, out=out)
    elif todo == 'ab,ac->cb':
        return gemm(B, A, out=out, trans1=True, trans2=True)
    elif todo == 'ab,bc->ac':
        return gemm(A, B, out=out, trans1=False, trans2=False)
    elif todo == 'ab,bc->ca':
        return gemm(B, A, out=out, trans1=True, trans2=True)
    elif todo == 'ab,ca->bc':
        return gemm(A, B, out=out, trans1=True, trans2=True)
    elif todo == 'ab,ca->cb':
        return gemm(B, A, out=out, trans1=False, trans2=False)
    elif todo == 'ab,cb->ac':
        return gemm(A, B, out=out, trans1=False, trans2=True)
    elif todo == 'ab,cb->ca':
        return gemm(B, A, out=out, trans1=False, trans2=True)
    else:
        raise ValueError("Invalid einsum: %s" % todo)