import torch
import sparseconvnet as scn
import matplotlib.pyplot as plt
import numpy as np


# Works if use_gpu=False, results in zeros if True
use_gpu = False


# Define model and input:
model = scn.Sequential().add(scn.SubmanifoldConvolution(2, 1, 1, 3, False)) \
                        .add(scn.SparseToDense(2, 1))

input = scn.InputBatch(2, [80, 80])

msg = [
    " X   X  XXX  X    X    XX     X       X   XX   XXX   X    XXX   ",
    " X   X  X    X    X   X  X    X       X  X  X  X  X  X    X  X  ",
    " XXXXX  XX   X    X   X  X    X   X   X  X  X  XXX   X    X   X ",
    " X   X  X    X    X   X  X     X X X X   X  X  X  X  X    X  X  ",
    " X   X  XXX  XXX  XXX  XX       X   X     XX   X  X  XXX  XXX   "]

# Make input sparse:
input.add_sample()
for y, line in enumerate(msg):
    for x, c in enumerate(line):
        if c == 'X':
            location = torch.LongTensor([x, y])
            featureVector = torch.FloatTensor([1])
            input.set_location(location, featureVector, 0)

# Mode to cpu or cuda:
if use_gpu:
    model.cuda()
    input.cuda()

# Calc result:
output = model.forward(input)
output_np = output.detach().cpu().numpy()[0].squeeze()

# Display Result:
print(output_np)
fig = plt.figure()
plt.imshow(output_np)
plt.show()