import torch
from torch.autograd import Variable
# from torch.nn.modules.module import Module

import os
import sys

sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname(__file__))))
import iisignature

from iisignature_torch import Sig, LogSig, SigJoin, SigScale


def trySig():
    inp = Variable(torch.randn(8, 2), requires_grad=True)
    result = Sig(inp, 2)
    print(result.data)
    result.backward(torch.randn(result.size()))
    print(inp.grad)


trySig()


def tryLogSig():
    s = iisignature.prepare(2, 2)
    for expanded in (False, True):
        inp = Variable(torch.randn(8, 2), requires_grad=True)
        if expanded:
            result = LogSig(inp, s, "x")
        else:
            result = LogSig(inp, s)
        print(result.data.numpy())
        result.backward(torch.randn(result.size()))
        print(inp.grad.data)


tryLogSig()


def trySigScale():
    inp1 = Variable(torch.randn(12, 6), requires_grad=True)
    inp2 = Variable(torch.randn(12, 2), requires_grad=True)
    result = SigScale(inp1, inp2, 2)
    print(result.data)
    result.backward(torch.randn(result.size()))
    print(inp1.grad)
    print(inp2.grad)


trySigScale()


def trySigJoin():
    inp1 = Variable(torch.randn(12, 6), requires_grad=True)
    inp2 = Variable(torch.randn(12, 2), requires_grad=True)
    result = SigJoin(inp1, inp2, 2)
    print(result.data)
    result.backward(torch.randn(result.size()))
    print(inp1.grad)
    print(inp2.grad)


trySigJoin()


def trySigJoin2():
    inp1 = Variable(torch.randn(12, 12), requires_grad=True)
    inp2 = Variable(torch.randn(12, 2), requires_grad=True)
    inp3 = Variable(torch.randn(1), requires_grad=True)
    result = SigJoin(inp1, inp2, 2, inp3)
    print(result.data)
    result.backward(torch.randn(result.size()))
    print(inp1.grad)
    print(inp2.grad)
    print(inp3.grad)


trySigJoin2()
