import io
import pickle

import numpy as np

import megengine as mge
import megengine.functional as F
import megengine.module as M
import megengine.utils.comp_graph_tools as cgtools
from megengine.core._trace_option import set_symbolic_shape
from megengine.jit import trace
from megengine.traced_module import trace_module

set_symbolic_shape(True)


class Main(M.Module):
    def forward(self, x):
        return x["data"]


class PreProcess(M.Module):
    def __init__(self):
        super().__init__()
        self.A = F.zeros((1,))
        self.I = F.ones((1,))
        self.bb_out = mge.tensor(
            np.array([[[0, 0], [160, 0], [160, 48], [0, 48]]], dtype="float32")
        )

    def forward(self, data, quad):
        """
        data: (1, 3, 48, 160)
        quad: (1, 4, 2)
        """
        N = quad.shape[0]
        dst = F.repeat(self.bb_out, N, axis=0).reshape(-1, 4, 2)
        I = F.broadcast_to(self.I, quad.shape)
        A = 
F.broadcast_to(self.A, (N, 8, 8))