import megengine.module as M
import megengine.functional as F
from megengine import amp

from .update import BasicUpdateBlock
from .extractor import BasicEncoder
from .corr import AGCL

from .attention import PositionEncodingSine, LocalFeatureTransformer


class CREStereo(M.Module):
    def __init__(self, max_disp=192, mixed_precision=False, test_mode=False):
        super(CREStereo, self).__init__()

        self.max_flow = max_disp
        self.mixed_precision = mixed_precision
        self.test_mode = test_mode

        self.hidden_dim = 128
        self.context_dim = 128
        self.dropout = 0

        # feature network and update block
        self.fnet = BasicEncoder(
            output_dim=256, norm_fn="instance", dropout=self.dropout
        )
        self.update_block = BasicUpdateBlock(
            hidden_dim=self.hidden_dim, cor_planes=4 * 9, mask_size=4
        )

        # loftr
        self.self_att_fn = LocalFeatureTransformer(
            d_model=256, nhead=8, layer_names=["self"] * 1, attention="linear"
        )
        self.cross_att_fn = LocalFeatureTransformer(
            d_model=256, nhead=8, layer_names=["cross"] * 1, attention="linear"
        )

        # adaptive search
        self.search_num = 9
        self.conv_offset_16 = M.Conv2d(
            256, self.search_num * 2, kernel_size=3, stride=1, padding=1
        )
        self.conv_offset_8 = M.Conv2d(
            256, self.search_num * 2, kernel_size=3, stride=1, padding=1
        )
        self.range_16 = 1
        self.range_8 = 1

    def freeze_bn(self):
        for m in self.modules():
            if isinstance(m, M.BatchNorm2d):
                m.eval()

    def unfold(self, x, kernel_size, dilation=1, padding=0, stride=1):
        n, c, h, w = x.shape
        if isinstance(kernel_size, tuple) or isinstance(kernel_size, list):
            assert len(kernel_size) == 2
            k1, k2 = kernel_size
        else:
            assert isinstance(kernel_size, int)
            k1 = k2 = kernel_size
        x = F.sliding_window(
            x,
            kernel_size=kernel_size,
            dilation=dilation,
            padding=padding,
            stride=stride,
        )
        x = F.reshape(x, (n, c, -1, k1 * k2))
        x = F.transpose(x, (0, 1, 3, 2))
        x = F.reshape(x, (n, c * k1 * k2, -1))
        return x

    def convex_upsample(self, flow, mask, rate=4):
        """[H/rate, W/rate, 2] -> [H, W, 2]"""
        N, _, H, W = flow.shape
        mask = F.reshape(mask, (N, 1, 9, rate, rate, H, W))
        mask = F.softmax(mask, axis=2)

        up_flow = self.unfold(rate * flow, [3, 3], padding=1)
        up_flow = 
F.reshape(up_flow, (N, 2, 9, 1, 1, H, W))