import vne
from vne.constants import nfeat, nsensors
from vne.model import simpleModel, init_weights_simple
from vne.persist import save_model, load_model
import numpy as np
import torch
import os


def encode_save(sig=np.random.random([1, nfeat, nsensors]), name='simpleModel_ini', dir_path="../src/vne/models"):
    """
    This function will create a model, test it, then save a persistent version (a file)

    Parameters
    __________
    sig:
        a numpy array with shape (nsamples, nfeatures, nsensors). in general a single neural signal may contain multiple channels.
        the multi-channel nature of the neural activations is a feature of the vne.
        typically shaped with size (1,1,S) where S is number os sensors
        the vne will map all the different channels to a single scalar encoded signal.

    name:
        string with the filename to save the model under
    dir:
        dir_path, the local directory to save the model

    Returns
    --------
    model:
        A copy of the encoder model generated by the function
    """

    model = simpleModel().eval()
    model.apply(init_weights_simple)
    sig = torch.tensor(sig.astype(np.float32)).to('cpu')
    enc = model(sig)
    print("signal={}".format(sig))
    print("encoded={}".format(enc))

    # save the model
    model.apply(vne.init_weights_simple)
    save_model(encoder=model, name=name, dir_path=dir_path)
    return model


def encode_load(sig=np.random.random([1, nfeat, nsensors]), name="simpleModel_ini", dir_path="../src/vne/models"):
    """
    This function will load a saved model, test it, then save a persistent version (a file)
    Parameters
    ----------
    sig:
        a numpy array with shape (nsamples, nfeatures, nsensors). in general a single neural signal may contain multiple channels.
        the multi-channel nature of the neural activations is a feature of the vne.
        typically shaped with size (1,1,S) where S is number os sensors
        the vne will map all the different channels to a single scalar encoded signal.

    name:
        the filename of the saved model

    dir_path:
        the directory path to the folder containing the file with the saved model
    """
    # load the saved model
    model = load_model(name, dir_path)

    # do some stuff

    # save the model
    model.apply(vne.init_weights_simple)
    return model


# Function to Convert to ONNX
def Convert_ONNX(model=None, name="simpleModel_ini", dir_path="../src/vne/models"):
    if model is None:
        model = encode_load()

    # set the model to inference mode (making sure)
    model.eval()
    name = os.path.join(dir_path, name)
    # Let's create a dummy input tensor
    dummy_input = torch.randn(1, nfeat, nsensors, requires_grad=True)

    # Export the model
    torch.onnx.export(model,  # model being run
                      dummy_input,  # model input (or a tuple for multiple inputs)
                      name + ".onnx",  # where to save the model
                      export_params=True,  # store the trained parameter weights inside the model file
                      opset_version=9,  # the ONNX version to export the model to
                      do_constant_folding=True,  # whether to execute constant folding for optimization
                      input_names=['sensorData'],  # the model's input names
                      output_names=['modelOutput'],  # the model's output names
                      dynamic_axes={'modelInput': {0: 'batch_size'},  # variable length axes
                                    'modelOutput': {0: 'batch_size'}})
    print(" ")
    print('Model has been converted to ONNX')


if __name__ == '__main__':
    # load persistent model from file
    model_name = 'simpleModel_ini-Trivial19'
    model = encode_save(name=model_name)

    print("saved model")
    # use the model
    sig = np.random.random([1, nfeat, nsensors])
    sig = torch.tensor(sig.astype(np.float32)).to('cpu')
    enc = model(sig)
    print("signal={}".format(sig))
    print("encoded={}".format(enc))

    print("ran model")
    Convert_ONNX(model, name=model_name)
