#!/usr/bin/python

# Copyright 2015-2016 Nervana Systems 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.
from __future__ import division
import numpy as np
#from ipdb import set_trace
from struct import pack, unpack
from neon import logger as neon_logger

def ceil_div(x, y):
    return -(-x // y)

def out_dim(S, X, padding, strides):
    return ceil_div(X - S + 1 + 2 * padding, strides)

def strip_mantissa(val):
    i = unpack('I', pack('f', val))[0] & 0x7f800000
    f = unpack('f', pack('I', i))[0]
    return f

def immediate(val):
    i = unpack('I', pack('f', val))[0] & 0x7ffff000
    f = unpack('f', pack('I', i))[0]
    return f

def quantize(ary, bits, sign=1):
    maxval = float(np.max(np.absolute(ary)))
    scale  = strip_mantissa(maxval) / float(1 << (bits - sign - 1))
    ary    = np.around(ary * (1.0 / scale)).astype(np.int64)
    return ary, np.float64(scale)

######### Direct Convolution #########

def fconv_slice(q, S, X, padding, strides):
    f1 = 0
    f2 = S - 1
    x1 = q * strides - padding
    x2 = x1 + f2
    if x1 < 0:
        f1 = -x1
        x1 = 0
    if x2 >= X:
        dif = x2 - X + 1
        f2 -= dif
        x2 -= dif
    return (slice(f1, f2 + 1), slice(x1, x2 + 1), f2 - f1 + 1)

def bconv_slice(x, S, Q, padding, strides):
    qs = x - (S - padding - 1)
    firstF = None
    for s in range(S): #TODO remove loop logic here.
        q = qs + s
        if q % strides == 0:
            q //= strides
            if q >= 0 and q < Q:
                if firstF is None:
                    firstF = s
                    firstE = q
                lastF = s
                lastE = q
    if firstF is None:
        return (slice(0,0,1), slice(0,0,1), 0)
    return (slice(firstF,lastF + 1,strides), slice(firstE,lastE + 1,1), 0)

def xprop_direct(I, F, O, padding, strides, backward=False):

    if all(x == 1 for x in F.shape[1:3]):
        C = F.shape[0]
        K = F.shape[4]
        if backward:
            # CxHWN = CxK . KxHWN
            O[:] = np.dot( F.reshape((C, -1)),   I.reshape((K, -1)) ).reshape((O.shape))
        else:
            # KxHWN = CxK.T . CxHWN
            O[:] = np.dot( F.reshape((C, -1)).T, I.reshape((C, -1)) ).reshape((O.shape))
        return

    if backward:
        # C <=> K and mirror R,S
        F = np.transpose(F[:,::-1,::-1,:], (3,1,2,0)).copy()
        xconv_slice = bconv_slice
    else:
        xconv_slice = fconv_slice

    C, Y, X, N = I.shape
    C, R, S, K = F.shape
    K, P, Q, N = O.shape

    qSlice = [ xconv_slice(q, S, X, padding[0], strides[0]) for q in range(Q) ]

    for p in range(P):
        sliceR, sliceY, _ = xconv_slice(p, R, Y, padding[1], strides[1])

        for q in range(Q):
            sliceS, sliceX, _ = qSlice[q]

            slicedF = F[:,sliceR,sliceS,:].reshape((-1, K))
            slicedI = I[:,sliceY,sliceX,:].reshape((-1, N))

            O[:,p,q,:] = np.dot( slicedF.T,  slicedI )

def updat_direct(I, E, U, padding, strides):

    C, Y, X, N = I.shape
    K, P, Q, N = E.shape
    C, R, S, K = U.shape

    if all(x == 1 for x in (R, S)):
        # CxK = CxHWN . KxHWN.T
        U[:] = np.dot( I.reshape((C, -1)), E.reshape((K, -1)).T ).reshape((U.shape))
        return

    U.fill(0.0)

    qSlice = [ fconv_slice(q, S, X, padding[0], strides[0]) for q in range(Q) ]

    for p in range(P):
        sliceR, sliceY, rlen = fconv_slice(p, R, Y, padding[1], strides[1])

        for q in range(Q):
            sliceS, sliceX, slen = qSlice[q]

            slicedI = I[:,sliceY,sliceX,:].reshape((-1, N))
            slicedE = E[:,p,q,:]

            U[:,sliceR,sliceS,:] += np.dot( slicedI,  slicedE.T ).reshape((C, rlen, slen, K))

######### Winograd Convolution #########



I_4x4_3x3 = (
    np.array([
    [ 4.0,  0.0, -5.0,  0.0, 1.0, 0.0],
    [ 0.0, -4.0, -4.0,  1.0, 1.0, 0.0],
    [ 0.0,  4.0, -4.0, -1.0, 1.0, 0.0],
    [ 0.0, -2.0, -1.0,  2.0, 1.0, 0.0],
    [ 0.0,  2.0, -1.0, -2.0, 1.0, 0.0],
    [ 0.0,  4.0,  0.0, -5.0, 0.0, 1.0]]),

    np.array([
    [ 4.0,  0.0, -5.0,  0.0, 1.0, 0.0],
    [ 0.0, -4.0, -4.0,  1.0, 1.0, 0.0],
    [ 0.0,  4.0, -4.0, -1.0, 1.0, 0.0],
    [ 0.0, -2.0, -1.0,  2.0, 1.0, 0.0],
    [ 0.0,  2.0, -1.0, -2.0, 1.0, 0.0],
    [ 0.0,  4.0,  0.0, -5.0, 0.0, 1.0]]),

    np.array([
    [ 1.0,  0.0,     -5.0 / 4.0,  0.0,      1.0 / 4.0, 0.0],
    [ 0.0,  2.0 / 3.0,  2.0 / 3.0, -1.0 / 6.0, -1.0 / 6.0, 0.0],
    [ 0.0, -2.0 / 3.0,  2.0 / 3.0,  1.0 / 6.0, -1.0 / 6.0, 0.0],
    [ 0.0, -1.0 / 12., -1.0 / 24.,  1.0 / 12.,  1.0 / 24., 0.0],
    [ 0.0,  1.0 / 12., -1.0 / 24., -1.0 / 12.,  1.0 / 24., 0.0],
    [ 0.0,  4.0,      0.0,     -5.0,      0.0,     1.0]]),

    np.array([
    [ 441.0 / 400.0,         0.0, -137.0 / 50.00,         0.0, 1.0, 0.0],
    [ 0.0,         -63.0 / 40.00,   -9.0 / 4.000,    7.0 / 10.0, 1.0, 0.0],
    [ 0.0,          63.0 / 40.00,   -9.0 / 4.000,   -7.0 / 10.0, 1.0, 0.0],
    [ 0.0,        -147.0 / 200.0,  -49.0 / 100.0,    3.0 / 2.00, 1.0, 0.0],
    [ 0.0,         147.0 / 200.0,  -49.0 / 100.0,   -3.0 / 2.00, 1.0, 0.0],
    [ 0.0,         441.0 / 400.0,          0.0, -137.0 / 50.0, 0.0, 1.0]]),

    np.array([
    [ 0.87890625,  0,          -2.640625,  0,        1, 0 ],
    [ 0,          -1.40625,    -2.25,      0.625,    1, 0 ],
    [ 0,           1.40625,    -2.25,     -0.625,    1, 0 ],
    [ 0,          -0.5859375,  -0.390625,  1.5,      1, 0 ],
    [ 0,           0.5859375,  -0.390625, -1.5,      1, 0 ],
    [ 0,           0.87890625,  0,        -2.640625, 0, 1 ]]),

    np.array([
    [ 0.87890625,  0,          -2.640625,  0,        1, 0 ],
    [ 0,          -1.40625,    -2.25,      0.625,    1, 0 ],
    [ 0,           1.40625,    -2.25,     -0.625,    1, 0 ],
    [ 0,          -0.5859375,  -0.390625,  1.5,      1, 0 ],
    [ 0,           0.5859375,  -0.390625, -1.5,      1, 0 ],
    [ 0,           0.87890625,  0,        -2.640625, 0, 1 ]]),
)

F_4x4_3x3 = (
    np.array([
    [  1.0 / 4.0,  0.0,      0.0     ],
    [ -1.0 / 6.0, -1.0 / 6.0, -1.0 / 6.0 ],
    [ -1.0 / 6.0,  1.0 / 6.0, -1.0 / 6.0 ],
    [  1.0 / 24.,  1.0 / 12.,  1.0 / 6.0 ],
    [  1.0 / 24., -1.0 / 12.,  1.0 / 6.0 ],
    [  0.0,      0.0,      1.0     ]]),

    np.array([
    [ 1.0,  0.0,  0.0 ],
    [ 1.0,  1.0,  1.0 ],
    [ 1.0, -1.0,  1.0 ],
    [ 1.0,  2.0,  4.0 ],
    [ 1.0, -2.0,  4.0 ],
    [ 0.0,  0.0,  1.0 ]]),

    np.array([
    [ 1.0,  0.0,  0.0 ],
    [ 1.0,  1.0,  1.0 ],
    [ 1.0, -1.0,  1.0 ],
    [ 1.0,  2.0,  4.0 ],
    [ 1.0, -2.0,  4.0 ],
    [ 0.0,  0.0,  1.0 ]]),

    np.array([
    [  400.0 / 441.00,          0.0,        0.0 ],
    [ -625.0 / 1078.0, -125.0 / 308.0, -25.0 / 88.0 ],
    [ -625.0 / 1078.0,  125.0 / 308.0, -25.0 / 88.0 ],
    [   25.0 / 198.00,   25.0 / 132.0,  25.0 / 88.0 ],
    [   25.0 / 198.00,  -25.0 / 132.0,  25.0 / 88.0 ],
    [           0.0,          0.0,        1.0 ]]),

    np.array([
    [ 1.13777777777778,   0,                  0                ],
    [-0.688403361344538, -0.430252100840336, -0.26890756302521 ],
    [-0.688403361344538,  0.430252100840336, -0.26890756302521 ],
    [ 0.119514472455649,  0.179271708683473,  0.26890756302521 ],
    [ 0.119514472455649, -0.179271708683473,  0.26890756302521 ],
    [ 0,                  0,                  1                ]]),

    np.array([
    [ 1, 1,         1 ,       1,     1,     0 ],
    [ 0, 0.625,    -0.625,    1.5,  -1.5,   0 ],
    [ 0, 0.390625,  0.390625, 2.25,  2.25,  1 ]]).T,
)

O_4x4_3x3 = (
    np.array([
    [ 1.0, 1.0,  1.0, 1.0,  1.0, 0.0 ],
    [ 0.0, 1.0, -1.0, 2.0, -2.0, 0.0 ],
    [ 0.0, 1.0,  1.0, 4.0,  4.0, 0.0 ],
    [ 0.0, 1.0, -1.0, 8.0, -8.0, 1.0 ]]),

    np.array([
    [ 1.0 / 4.0, -1.0 / 6.0, -1.0 / 6.0, 1.0 / 24.,  1.0 / 24., 0.0 ],
    [ 0.0,     -1.0 / 6.0,  1.0 / 6.0, 1.0 / 12., -1.0 / 12., 0.0 ],
    [ 0.0,     -1.0 / 6.0, -1.0 / 6.0, 1.0 / 6.0,  1.0 / 6.0, 0.0 ],
    [ 0.0,     -1.0 / 6.0,  1.0 / 6.0, 1.0 / 3.0, -1.0 / 3.0, 1.0 ]]),

    np.array([
    [ 1.0, 1.0,  1.0, 1.0,  1.0, 0.0 ],
    [ 0.0, 1.0, -1.0, 2.0, -2.0, 0.0 ],
    [ 0.0, 1.0,  1.0, 4.0,  4.0, 0.0 ],
    [ 0.0, 1.0, -1.0, 8.0, -8.0, 1.0 ]]),

    np.array([
    [ 1.0,       1.0000,        1.0000,      1.0,       1.0, 0.0 ],
    [ 0.0,   7.0 / 10.000,   -7.0 / 10.000,  3.0 / 2.0,  -3.0 / 2.0, 0.0 ],
    [ 0.0,  49.0 / 100.00,   49.0 / 100.00,  9.0 / 4.0,   9.0 / 4.0, 0.0 ],
    [ 0.0, 343.0 / 1000.0, -343.0 / 1000.0, 27.0 / 8.0, -27.0 / 8.0, 1.0 ]]),

    np.array([
    [ 1, 1,            1,           1,      1,     0 ],
    [ 0, 0.625,       -0.625,       1.5,   -1.5,   0 ],
    [ 0, 0.390625,     0.390625,    2.25,   2.25,  0 ],
    [ 0, 0.244140625, -0.244140625, 3.375, -3.375, 1 ]]),

    np.array([
    [ 1.13777777777778,   0,                  0,                 0,                ],
    [-0.688403361344538, -0.430252100840336, -0.26890756302521, -0.168067226890756 ],
    [-0.688403361344538,  0.430252100840336, -0.26890756302521,  0.168067226890756 ],
    [ 0.119514472455649,  0.179271708683473,  0.26890756302521,  0.403361344537815 ],
    [ 0.119514472455649, -0.179271708683473,  0.26890756302521, -0.403361344537815 ],
    [ 0,                  0,                  0,                 1,                ]]).T,
)

rcp3  = 1.0 / 3.0
rcp4  = 1.0 / 4.0
rcp6  = 1.0 / 6.0
rcp12 = 1.0 / 12.0
rcp24 = 1.0 / 24.0

def trans_I_4x4_3x3(Iw, I, minimal=False, trans=False):
    if minimal:

        T0 = np.empty((6,6))
        T1 = np.empty((6,6))