import torch
import torch.nn as nn
import onnxruntime
import numpy as np
import cv2
import random

# Set random seeds
seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.set_printoptions(precision=15)


class LayerLSTM(nn.Module):
    def __init__(self, features, hidden_size, layers=1, batch_first=False):
        super(LayerLSTM, self).__init__()
        self.rnns = nn.LSTM(
            input_size=features,
            hidden_size=hidden_size,
            num_layers=layers,
            batch_first=batch_first
        )
        self.set_weights_to_ones()

    def forward(self, x, hx, cx):
        x, (hx, cx) = self.rnns(x, (hx,cx))
        return x, (hx, cx)

    def set_weights_to_ones(self):
        with torch.no_grad():
            for name, param in self.rnns.named_parameters():
                if 'weight' in name or 'bias' in name:
                    param.fill_(1.0)

if __name__ == "__main__":

    features    = 3
    hidden_size = 7
    batch_size  = 5
    seq_len     = 2
    layout      = True

    model = LayerLSTM(features, hidden_size, batch_first=layout)

    with torch.no_grad():
        if layout:
            x  = torch.ones(batch_size, seq_len, features)
        else:
            x  = torch.ones(seq_len, batch_size, features)

        hx = torch.ones(1, batch_size, hidden_size)
        cx = torch.ones(1, batch_size, hidden_size)

        print(x.shape, hx.shape, cx.shape, sep='\n')
        out, (hn, cn) = model(x, hx, cx)
        
    torch.onnx.export(
       model,
       (x, hx, cx),
       './sample_lstm.onnx',
       verbose=True,
       input_names=['x', 'hx', 'cx'],
       output_names=['output']
    )
    
    net = cv2.dnn.readNetFromONNX('./sample_lstm.onnx')

    x_opencv  = x.numpy()[..., None]
    hx_opencv = hx.numpy()[..., None]
    cx_opencv = cx.numpy()[..., None]

    # Set input data
    net.setInput(x_opencv)
    net.setInput(hx_opencv, name="hx")
    net.setInput(cx_opencv, name="cx")
    
    out_opencv = net.forward()