import torch
from torch import nn
from torch.cuda.amp import autocast
from rscd.models.decoderheads.vision_lstm import ViLBlock, SequenceTraversal
from torch.nn import functional as F
from functools import partial
from rscd.models.backbones.lib_mamba.vmambanew import SS2D
import pywt

class PA(nn.Module):
    def __init__(self, dim, norm_layer, act_layer):
        super().__init__()
        self.p_conv = nn.Sequential(
            nn.Conv2d(dim, dim*4, 1, bias=False),
            norm_layer(dim*4),
            act_layer(),
            nn.Conv2d(dim*4, dim, 1, bias=False)
        )
        self.gate_fn = nn.Sigmoid()

    def forward(self, x):
        att = self.p_conv(x)
        x = x * self.gate_fn(att)

        return x

class Mish(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x * torch.tanh(F.softplus(x))


class _ScaleModule(nn.Module):
    def __init__(self, dims, init_scale=1.0):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(*dims) * init_scale)

    def forward(self, x):
        return torch.mul(self.weight, x)

def create_wavelet_filter(wave, in_size, out_size, dtype=torch.float):
    w = pywt.Wavelet(wave)
    dec_hi = torch.tensor(w.dec_hi[::-1], dtype=dtype)
    dec_lo = torch.tensor(w.dec_lo[::-1], dtype=dtype)
    dec_filters = torch.stack([dec_lo.unsqueeze(0) * dec_lo.unsqueeze(1),
                               dec_lo.unsqueeze(0) * dec_hi.unsqueeze(1),
                               dec_hi.unsqueeze(0) * dec_lo.unsqueeze(1),
                               dec_hi.unsqueeze(0) * dec_hi.unsqueeze(1)], dim=0)
    dec_filters = dec_filters[:, None].repeat(in_size, 1, 1, 1)
    rec_hi = torch.tensor(w.rec_hi[::-1], dtype=dtype).flip(dims=[0])
    rec_lo = torch.tensor(w.rec_lo[::-1], dtype=dtype).flip(dims=[0])
    rec_filters = torch.stack([rec_lo.unsqueeze(0) * rec_lo.unsqueeze(1),
                               rec_lo.unsqueeze(0) * rec_hi.unsqueeze(1),
                               rec_hi.unsqueeze(0) * rec_lo.unsqueeze(1),
                               rec_hi.unsqueeze(0) * rec_hi.unsqueeze(1)], dim=0)
    rec_filters = rec_filters[:, None].repeat(out_size, 1, 1, 1)
    return dec_filters, rec_filters

def wavelet_transform(x, filters):
    b, c, h, w = x.shape
    pad = (filters.shape[2] // 2 - 1, filters.shape[3] // 2 - 1)
    x = F.conv2d(x, filters, stride=2, groups=c, padding=pad)
    x = x.reshape(b, c, 4, h // 2, w // 2)
    return x

def inverse_wavelet_transform(x, filters):
    b, c, _, h_half, w_half = x.shape
    pad = (filters.shape[2] // 2 - 1, filters.shape[3] // 2 - 1)
    x = x.reshape(b, c * 4, h_half, w_half)
    x = F.conv_transpose2d(x, filters, stride=2, groups=c, padding=pad)
    return x

class MBWTConv2d(nn.Module):
    def __init__(self, in_channels, kernel_size=5, wt_levels=1, wt_type='db1', ssm_ratio=1, forward_type="v05"):
        super().__init__()
        assert in_channels == in_channels
        self.wt_levels = wt_levels
        self.wt_filter, self.iwt_filter = create_wavelet_filter(wt_type, in_channels, in_channels)
        self.wt_filter = nn.Parameter(self.wt_filter, requires_grad=False)
        self.iwt_filter = nn.Parameter(self.iwt_filter, requires_grad=False)
        self.wt_function = partial(wavelet_transform, filters=self.wt_filter)
        self.iwt_function = partial(inverse_wavelet_transform, filters=self.iwt_filter)
        self.global_atten = SS2D(d_model=in_channels, d_state=1, ssm_ratio=ssm_ratio, initialize="v2",
                                 forward_type=forward_type, channel_first=True, k_group=2)
        self.base_scale = _ScaleModule([1, in_channels, 1, 1])
        self.wavelet_convs = nn.ModuleList([
            nn.Conv2d(in_channels * 4, in_channels * 4, kernel_size, padding='same', groups=in_channels * 4)
            for _ in range(wt_levels)
        ])
        self.wavelet_scale = nn.ModuleList([
            _ScaleModule([1, in_channels * 4, 1, 1], init_scale=0.1)
            for _ in range(wt_levels)
        ])

    def forward(self, x):
        x_ll_in_levels, x_h_in_levels, shapes_in_levels = [], [], []
        curr_x_ll = x
        for i in range(self.wt_levels):
            curr_shape = curr_x_ll.shape
            shapes_in_levels.append(curr_shape)
            if (curr_shape[2] % 2 > 0) or (curr_shape[3] % 2 > 0):
                curr_x_ll = F.pad(curr_x_ll, (0, curr_shape[3] % 2, 0, curr_shape[2] % 2))
            curr_x = self.wt_function(curr_x_ll)
            curr_x_ll = curr_x[:, :, 0, :, :]
            shape_x = curr_x.shape
            curr_x_tag = curr_x.reshape(shape_x[0], shape_x[1] * 4, shape_x[3], shape_x[4])
            curr_x_tag = self.wavelet_scale[i](self.wavelet_convs[i](curr_x_tag)).reshape(shape_x)
            x_ll_in_levels.append(curr_x_tag[:, :, 0, :, :])
            x_h_in_levels.append(curr_x_tag[:, :, 1:4, :, :])
        next_x_ll = 0
        for i in range(self.wt_levels - 1, -1, -1):
            curr_x_ll = x_ll_in_levels.pop() + next_x_ll
            curr_x = torch.cat([curr_x_ll.unsqueeze(2), x_h_in_levels.pop()], dim=2)
            next_x_ll = self.iwt_function(curr_x)
            next_x_ll = next_x_ll[:, :, :shapes_in_levels[i][2], :shapes_in_levels[i][3]]
        x_tag = next_x_ll
        x = self.base_scale(self.global_atten(x)) + x_tag
        return x

class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.fc1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
        out = avg_out + max_out
        return self.sigmoid(out)

# 空间注意力模块
class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()

        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1

        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x = torch.cat([avg_out, max_out], dim=1)
        x = self.conv1(x)
        return self.sigmoid(x)

# CBAM 注意力模块
class CBAM(nn.Module):
    def __init__(self, in_planes):
        super(CBAM, self).__init__()
        self.ca = ChannelAttention(in_planes)
        self.sa = SpatialAttention()

    def forward(self, x):
        x = self.ca(x) * x
        x = self.sa(x) * x
        return x

class DynamicConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, groups=1, bias=False, num_experts=4):
        super(DynamicConv2d, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.groups = groups
        self.bias = bias
        self.num_experts = num_experts

        self.experts = nn.ModuleList([
            nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, groups=groups, bias=bias)
            for _ in range(num_experts)
        ])
        self.gating = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, num_experts, 1, bias=False),
            nn.Softmax(dim=1)
        )

    def forward(self, x):
        gates = self.gating(x)
        gates = gates.view(x.size(0), self.num_experts, 1, 1, 1)
        outputs = []
        for i, expert in enumerate(self.experts):
            outputs.append(expert(x).unsqueeze(1))
        outputs = torch.cat(outputs, dim=1)
        out = (gates * outputs).sum(dim=1)
        return out

class DWConv2d_BN_ReLU(nn.Sequential):
    def __init__(self, in_channels, out_channels, kernel_size=3):
        super().__init__()
        self.add_module('dwconv3x3', DynamicConv2d(in_channels, in_channels, kernel_size=kernel_size,
                                                   stride=1, padding=kernel_size // 2, groups=in_channels, bias=False))
        self.add_module('bn1', nn.BatchNorm2d(in_channels))
        self.add_module('relu', Mish())
        self.add_module('dwconv1x1', nn.Conv2d(in_channels, out_channels, kernel_size=1,
                                               stride=1, padding=0, groups=in_channels, bias=False))
        self.add_module('bn2', nn.BatchNorm2d(out_channels))

class Conv2d_BN(nn.Sequential):
    def __init__(self, a, b, ks=1, stride=1, pad=0, groups=1):
        super().__init__()
        self.add_module('c', nn.Conv2d(a, b, ks, stride, pad, groups=groups, bias=False))
        self.add_module('bn', nn.BatchNorm2d(b))

class FFN(nn.Module):
    def __init__(self, ed, h):
        super().__init__()
        self.pw1 = Conv2d_BN(ed, h)
        self.act = Mish()
        self.pw2 = Conv2d_BN(h, ed)

    def forward(self, x):
        return self.pw2(self.act(self.pw1(x)))

class StochasticDepth(nn.Module):
    def __init__(self, survival_prob=0.8):
        super().__init__()
        self.survival_prob = survival_prob

    def forward(self, x):
        if not self.training:
            return x
        batch_size = x.shape[0]
        random_tensor = self.survival_prob + torch.rand([batch_size, 1, 1, 1], dtype=x.dtype, device=x.device)
        binary_tensor = torch.floor(random_tensor)
        return x * binary_tensor / self.survival_prob

class Residual(nn.Module):
    def __init__(self, m, survival_prob=0.8):
        super().__init__()
        self.m = m
        self.stochastic_depth = StochasticDepth(survival_prob)

    def forward(self, x):
        return x + self.stochastic_depth(self.m(x))

class GLP_block(nn.Module):
    def __init__(self, dim, global_ratio=0.25, local_ratio=0.25,pa_ratio = 0.1, kernels=3, ssm_ratio=1, forward_type="v052d"):
        super().__init__()
        self.dim = dim
        self.global_channels = int(global_ratio * dim)
        self.local_channels = int(local_ratio * dim)
        self.pa_channels = int(pa_ratio * dim)
        self.identity_channels = dim - self.global_channels - self.local_channels - self.pa_channels
        self.local_op = nn.ModuleList([
            DWConv2d_BN_ReLU(self.local_channels, self.local_channels, k)
            for k in [3, 5, 7]
        ]) if self.local_channels > 0 else nn.Identity()
        self.global_op = MBWTConv2d(self.global_channels, kernel_size=kernels,
                                     ssm_ratio=ssm_ratio, forward_type=forward_type) \
            if self.global_channels > 0 else nn.Identity()
        self.cbam = CBAM(dim)
        self.proj = nn.Sequential(
            Mish(),
            Conv2d_BN(dim, dim),
            CBAM(dim)
        )

        self.pa_op = PA(self.pa_channels, norm_layer=nn.BatchNorm2d, act_layer=nn.GELU) \
            if self.pa_channels > 0 else nn.Identity()

    def forward(self, x):
        x1, x2, x3, x4 = torch.split(x, [self.global_channels, self.local_channels, self.identity_channels, self.pa_channels], dim=1)
        if isinstance(self.local_op, nn.ModuleList):
            local_features = [op(x2) for op in self.local_op]
            local_features = torch.cat(local_features, dim=1)
            local_features = torch.mean(local_features, dim=1, keepdim=True)
            local_features = local_features.expand(-1, self.local_channels, -1, -1)
        else:
            local_features = self.local_op(x2)
        out = torch.cat([self.global_op(x1), local_features, x3, self.pa_op(x4)], dim=1)
        return self.proj(out)





class SASF(nn.Module):
    def __init__(self, dim, global_ratio=0.25, local_ratio=0.25,pa_ratio = 0.1, kernels=3, ssm_ratio=1, forward_type="v052d"):
        super().__init__()
        self.dim = dim
        self.global_channels = int(global_ratio * dim)
        self.local_channels = int(local_ratio * dim)
        self.pa_channels = int(pa_ratio * dim)
        self.identity_channels = dim - self.global_channels - self.local_channels - self.pa_channels
        self.local_op = nn.ModuleList([
            DWConv2d_BN_ReLU(self.local_channels, self.local_channels, k)
            for k in [3, 5, 7]
        ]) if self.local_channels > 0 else nn.Identity()
        self.global_op = MBWTConv2d(self.global_channels, kernel_size=kernels,
                                     ssm_ratio=ssm_ratio, forward_type=forward_type) \
            if self.global_channels > 0 else nn.Identity()
        self.cbam = CBAM(dim)
        self.proj = nn.Sequential(
            Mish(),
            Conv2d_BN(dim, dim),
            CBAM(dim)
        )

        self.pa_op = PA(self.pa_channels, norm_layer=nn.BatchNorm2d, act_layer=nn.GELU) \
            if self.pa_channels > 0 else nn.Identity()

    def forward(self, x):
        x1, x2, x3, x4 = torch.split(x, [self.global_channels, self.local_channels, self.identity_channels, self.pa_channels], dim=1)
        if isinstance(self.local_op, nn.ModuleList):
            local_features = [op(x2) for op in self.local_op]
            local_features = torch.cat(local_features, dim=1)
            local_features = torch.mean(local_features, dim=1, keepdim=True)
            local_features = local_features.expand(-1, self.local_channels, -1, -1)
        else:
            local_features = self.local_op(x2)
        out = torch.cat([self.global_op(x1), local_features, x3, self.pa_op(x4)], dim=1)
        return self.proj(out)


class ViLLayer(nn.Module):
    def __init__(self, dim, d_state = 16, d_conv = 4, expand = 2):
        super().__init__()
        self.dim = dim
        self.norm = nn.LayerNorm(dim)
        self.vil = ViLBlock(
            dim= self.dim,
            direction=SequenceTraversal.ROWWISE_FROM_TOP_LEFT
        )
    
    @autocast(enabled=False)
    def forward(self, x):
        if x.dtype == torch.float16:
            x = x.type(torch.float32)
        B, C = x.shape[:2]
        assert C == self.dim
        n_tokens = x.shape[2:].numel()
        img_dims = x.shape[2:]
        x_flat = x.reshape(B, C, n_tokens).transpose(-1, -2)
        x_vil = self.vil(x_flat)
        out = x_vil.transpose(-1, -2).reshape(B, C, *img_dims)

        return out

def dsconv_3x3(in_channel, out_channel):
    return nn.Sequential(
        nn.Conv2d(in_channel, in_channel, kernel_size=3, stride=1, padding=1, groups=in_channel),
        nn.Conv2d(in_channel, out_channel, kernel_size=1, stride=1, padding=0, groups=1),
        nn.BatchNorm2d(out_channel),
        nn.ReLU(inplace=True)
    )

def conv_1x1(in_channel, out_channel):
    return nn.Sequential(
        nn.Conv2d(in_channel, out_channel, kernel_size=1, stride=1, padding=0, bias=False),
        nn.BatchNorm2d(out_channel),
        nn.ReLU(inplace=True)
    )

class SqueezeAxialPositionalEmbedding(nn.Module):
    def __init__(self, dim, shape):
        super().__init__()

        self.pos_embed = nn.Parameter(torch.randn([1, dim, shape]))

    def forward(self, x):
        B, C, N = x.shape
        x = x + F.interpolate(self.pos_embed, size=(N), mode='linear', align_corners=False)

        return x

class SEBlock(nn.Module):
    def __init__(self, channels, r=16):
        super().__init__()
        self.fc = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(channels, channels//r, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(channels//r, channels, 1),
            nn.Sigmoid()
        )
    def forward(self, x):
        w = self.fc(x)  # (B, C, 1, 1)
        return x * w
class CTTF1(nn.Module):
    def __init__(self, in_channel, out_channel,global_ratio=0.2, local_ratio=0.2, pa_ratio = 0.2 ,kernels=5, ssm_ratio=2.0, forward_type="v052d"):
        super().__init__()
        self.catconvA = dsconv_3x3(in_channel * 2, in_channel)
        self.catconvB = dsconv_3x3(in_channel * 2, in_channel)
        self.catconv  = dsconv_3x3(in_channel * 2, out_channel)
        self.convA    = nn.Conv2d(in_channel, 1, 1)
        self.convB    = nn.Conv2d(in_channel, 1, 1)
        self.sigmoid  = nn.Sigmoid()

        self.mixer = Residual(GLP_block(in_channel, global_ratio, local_ratio,pa_ratio, kernels, ssm_ratio, forward_type))
        self.mixer2 = Residual(
            SASF(in_channel, global_ratio = 0, local_ratio = 0.1, pa_ratio = 0, kernels = 5, ssm_ratio = 1, forward_type = "v052d"))

        self.fuse = nn.Sequential(
        nn.Conv2d(in_channel * 3, in_channel, kernel_size=1),
        nn.ReLU(inplace=True)
        )
        self.cbam = CBAM(in_channel * 3)

        self.act = nn.SiLU()
    def forward(self, xA, xB):
        x_diffA = self.mixer(xA)
        x_diffB = self.mixer(xB)

        f1 = x_diffA
        f2 = x_diffB
        diff_signed = f1 - f2
        diff_abs = torch.abs(diff_signed)
        sum_feat = f1 + f2

        diff_signed = self.mixer2(diff_signed)
        diff_abs = self.mixer2(diff_abs)
        sum_feat = self.mixer2(sum_feat)
        # 将多路特征在通道维度拼接
        f_fuse = torch.cat([diff_signed, diff_abs, sum_feat], dim=1)  # (B, 4C, H, W)
        # 再接一个 1x1 卷积降维或提炼信息
        f_fuse = self.cbam(f_fuse)
        x_diff = self.fuse(f_fuse)

        return x_diff

class CTTF2(nn.Module):
    def __init__(self, in_channel, out_channel, global_ratio=0.25, local_ratio=0.25, pa_ratio=0, kernels=7,
                 ssm_ratio=2.0, forward_type="v052d"):
        super().__init__()
        self.catconvA = dsconv_3x3(in_channel * 2, in_channel)
        self.catconvB = dsconv_3x3(in_channel * 2, in_channel)
        self.catconv = dsconv_3x3(in_channel * 2, out_channel)
        self.convA = nn.Conv2d(in_channel, 1, 1)
        self.convB = nn.Conv2d(in_channel, 1, 1)
        self.sigmoid = nn.Sigmoid()


        self.mixer = Residual(
            GLP_block(in_channel, global_ratio, local_ratio, pa_ratio, kernels, ssm_ratio, forward_type))
        self.mixer2 = Residual(
            SASF(in_channel, global_ratio=0, local_ratio=0.1, pa_ratio=0, kernels=5, ssm_ratio=1,
                              forward_type="v052d"))

        self.fuse = nn.Sequential(
            nn.Conv2d(in_channel * 3, in_channel, kernel_size=1),
            nn.ReLU(inplace=True)
        )
        self.cbam = CBAM(in_channel * 3)

        self.act = nn.SiLU()

    def forward(self, xA, xB):
        x_diffA = self.mixer(xA)
        x_diffB = self.mixer(xB)

        f1 = x_diffA
        f2 = x_diffB
        diff_signed = f1 - f2
        diff_abs = torch.abs(diff_signed)
        sum_feat = f1 + f2

        diff_signed = self.mixer2(diff_signed)
        diff_abs = self.mixer2(diff_abs)
        sum_feat = self.mixer2(sum_feat)
        f_fuse = torch.cat([diff_signed, diff_abs, sum_feat], dim=1)  # (B, 4C, H, W)
        f_fuse = self.cbam(f_fuse)
        x_diff = self.fuse(f_fuse)

        return x_diff

class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.,channels_first=True):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features

        Linear = partial(nn.Conv2d, kernel_size=1, padding=0) if channels_first else nn.Linear
        self.fc1 = Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x
    
class LHBlock(nn.Module):
    def __init__(self, channels_l, channels_h):
        super().__init__()
        self.channels_l = channels_l
        self.channels_h = channels_h
        self.cross_size = 12
        self.cross_kv = nn.Sequential(
            nn.BatchNorm2d(channels_l),
            nn.AdaptiveMaxPool2d(output_size=(self.cross_size, self.cross_size)),
            nn.Conv2d(channels_l, 2 * channels_h, 1, 1, 0)
        )

        self.conv = conv_1x1(channels_l, channels_h)
        self.norm = nn.BatchNorm2d(channels_h)
        
        self.mlp_l = Mlp(in_features=channels_l, out_features=channels_l)
        self.mlp_h = Mlp(in_features=channels_h, out_features=channels_h)

    def _act_sn(self, x):
        _, _, H, W = x.shape
        inner_channel = self.cross_size * self.cross_size
        x = x.reshape([-1, inner_channel, H, W]) * (inner_channel**-0.5)
        x = F.softmax(x, dim=1)
        x = x.reshape([1, -1, H, W])
        return x
    
    def attn_h(self, x_h, cross_k, cross_v):
        B, _, H, W = x_h.shape
        x_h = self.norm(x_h)
        x_h = x_h.reshape([1, -1, H, W])  # n,c_in,h,w -> 1,n*c_in,h,w
        x_h = F.conv2d(x_h, cross_k, bias=None, stride=1, padding=0,
                        groups=B)  # 1,n*c_in,h,w -> 1,n*144,h,w  (group=B)
        x_h = self._act_sn(x_h)
        x_h = F.conv2d(x_h, cross_v, bias=None, stride=1, padding=0,
                        groups=B)  # 1,n*144,h,w -> 1, n*c_in,h,w  (group=B)
        x_h = x_h.reshape([-1, self.channels_h, H,
                        W])  # 1, n*c_in,h,w -> n,c_in,h,w  (c_in = c_out)

        return x_h

    def forward(self, x_l, x_h):
        x_l = x_l + self.mlp_l(x_l)
        x_l_conv = self.conv(x_l)
        x_h = x_h + F.interpolate(x_l_conv, size=x_h.shape[2:], mode='bilinear')

        cross_kv = self.cross_kv(x_l)
        cross_k, cross_v = cross_kv.split(self.channels_h, 1)
        cross_k = cross_k.permute(0, 2, 3, 1).reshape([-1, self.channels_h, 1, 1])  # n*144,channels_h,1,1
        cross_v = cross_v.reshape([-1, self.cross_size * self.cross_size, 1, 1])  # n*channels_h,144,1,1

        x_h = x_h + self.attn_h(x_h, cross_k, cross_v) # [4, 40, 128, 128]
        x_h = x_h + self.mlp_h(x_h)

        return x_h


class CTTF(nn.Module):
    def __init__(self, channels=[40, 80, 192, 384]):
        super().__init__()
        self.channels = channels
        self.fusion0 = CTTF1(channels[0], channels[0])
        self.fusion1 = CTTF1(channels[1], channels[1])
        self.fusion2 = CTTF2(channels[2], channels[2])
        self.fusion3 = CTTF2(channels[3], channels[3])

        self.LHBlock1 = LHBlock(channels[1], channels[0])
        self.LHBlock2 = LHBlock(channels[2], channels[0])
        self.LHBlock3 = LHBlock(channels[3], channels[0])

        self.mlp1 = Mlp(in_features=channels[0], out_features=channels[0])
        self.mlp2 = Mlp(in_features=channels[0], out_features=2)
        self.dwc = dsconv_3x3(channels[0], channels[0])

    def forward(self, inputs):
        featuresA, featuresB = inputs
        # fA_0, fA_1, fA_2, fA_3 = featuresA
        # fB_0, fB_1, fB_2, fB_3 = featuresB
        x_diff_0 = self.fusion0(featuresA[0], featuresB[0]) # [4, 40, 128, 128]
        x_diff_1 = self.fusion1(featuresA[1], featuresB[1]) # [4, 80, 64, 64]
        # x_diff_2 = featuresA[2] -  featuresB[2]
        # x_diff_3 = featuresA[3] -  featuresB[3]
        x_diff_2 = self.fusion2(featuresA[2], featuresB[2]) # [4, 192, 32, 32]
        x_diff_3 = self.fusion3(featuresA[3], featuresB[3]) # [4, 384, 16, 16]
        
        x_h = x_diff_0
        x_h = self.LHBlock1(x_diff_1, x_h) # [4, 40, 128, 128]
        x_h = self.LHBlock2(x_diff_2, x_h)
        x_h = self.LHBlock3(x_diff_3, x_h)

        out = self.mlp2(self.dwc(x_h) + self.mlp1(x_h))

        out = F.interpolate(
            out,
            scale_factor=(4, 4),
            mode="bilinear",
            align_corners=False,
        )
        return out