#!/usr/bin/env python3

import nnabla as nn
import nnabla.functions as F  # it crashes without this
import numpy.random as R
import itertools as IT

# I hate you, Python
import os
from pathlib import Path
exec(Path(os.path.dirname(__file__) + '/' + 'ntshared.py').read_text())

seed()

x, t, y, loss, hs = setup_network()

train_network(loss)

R.seed() # reseed for test data

pq, label = random_data()
preds, loss = predict(pq, label)

for name, param in nn.get_parameters().items():
    print(name, param.shape, param.g.flat[:20])

print("Test loss:", loss.d)

#plot_classified(x.d, t.d.reshape(BATCH_SIZE), preds)

def nnabla_to_smt2(var, names={}, collect={}, rcollect={}, assertions=[],
                   nid=0, normal=True):
    if var in rcollect:
        return collect  # already processed this variable
    rcollect[var] = nid
    if var not in names:
        names[var] = 'var_{}'.format(nid)
    collect[nid] = var
    cur_name = names[var]
    nid += 1
    if var.parent is not None:
        print(var.parent)
        print(var.parent.inputs)
        print(type(var.parent.inputs))
        for index, input in enumerate(var.parent.inputs):
            _, _, nid = nnabla_to_smt2(input, names, collect, rcollect,
                                       assertions, nid, index == 0)

        if var.parent.name == 'ReLU':
            assert normal
            assert len(var.parent.inputs) == 1
            assert len(var.shape) == 2
            assert var.parent.inputs[0].shape == var.shape
            param_name = names[var.parent.inputs[0]]
            for index in range(var.shape[1]):
                assertions.append('(= {}_{} (max 0 {}_{}))'.format(
                    cur_name, index, param_name, index
                ))
        elif var.parent.name == 'Affine':
            # Wx + b -- W and b are trained parameters
            assert normal
            assert len(var.shape) == 2
            assert len(var.parent.inputs) == 3
            var_x = var.parent.inputs[0]
            var_W = var.parent.inputs[1]
            var_b = var.parent.inputs[2]
            assert len(var_x.shape) == 2
            assert len(var_W.shape) == 2
            assert len(var_b.shape) == 1
            assert var_W.shape[0] == var_x.shape[1]
            assert var_W.shape[1] == var.shape[1]
            assert var_W.shape[1] == var_b.shape[0]
            x_name = names[var_x]
            for i in range(var.shape[1]):
                terms = []
                for j in range(var_x.shape[1]):
                    terms.append('(* {} {}_{})'.format(
                        var_W.d[j][i],
                        x_name,
                        j
                    ))
                assertions.append('(= {}_{} (+ {} {}))'.format(
                    cur_name,
                    i,
                    var_b.d[i],
                    ' '.join(terms)
                ))
        else:
            raise Exception('Unsupported function: {}'.format(var.parent.name))
    return collect, assertions, nid

collect, assertions, _ = nnabla_to_smt2(y, {x: 'x', y: 'y'})
for nid, var in collect.items():
    print(nid, var, var.ndim)
for assertion in assertions:
    print(assertion)

@@ -629,16 +629,21 @@
 _to_smt2
+_info
 (var, na
@@ -690,16 +690,21 @@
 ons=%5B%5D,%0A
+     
         
@@ -1195,16 +1195,21 @@
 _to_smt2
+_info
 (input, 
@@ -1234,16 +1234,21 @@
 ollect,%0A
+     
         
@@ -3033,178 +3033,260 @@
 id%0A%0A
-collect, assertions, _ = nnabla_to_smt2(y, %7Bx: 'x', y: 'y'%7D)%0Afor nid, var in collect.items():%0A    print(nid, var, var.ndim)%0Afor assertion in assertions:%0A    print(assertion
+def nnabla_to_smt2(var, names=%7B%7D):%0A    collect, assertions, _ = nnabla_to_smt2_info(var, names)%0A    smt2 = ''%0A    smt2 += ''.join(map(lambda a: '(assert %7B%7D)%5Cn'.format(a), assertions))%0A    return smt2%0A%0Asmt2 = nnabla_to_smt2(y, %7Bx: 'x', y: 'y'%7D)%0Aprint(smt2
 )%0A

Collect smt2 into string.