from PepperPepper.environment import torch, nn, math, trunc_normal_, profile
from PepperPepper.IRSTD.models import SCTransNet, MiM, get_SCTrans_config



class IRSTDNet(nn.Module):
    def __init__(self, model_name):
        super(IRSTDNet, self).__init__()
        self.model_name = model_name
        self.cal_loss = nn.BCEWithLogitsLoss()
        self.model = None


        if model_name == 'MiM':
            self.model = MiM()
        elif model_name == 'SCTransNet':
            config = get_SCTrans_config()
            self.model = SCTransNet(config, mode='train', deepsuper=True)
        else:
            raise NotImplementedError

        self.apply(self._init_weights)

    def forward(self, img):
        return self.model(img)

    def loss(self, preds, gt_masks):
        # preds = torch.sigmoid(preds)
        # gt_masks = torch.sigmoid(gt_masks)
        if isinstance(preds, list):
            loss_total = 0
            for i in range(len(preds)):
                pred = preds[i]
                gt_mask = gt_masks[i]
                loss = self.cal_loss(pred, gt_mask)
                loss_total = loss_total + loss
            return loss_total / len(preds)

        elif isinstance(preds, tuple):
            loss_total = 0
            for i in range(len(preds)):
                pred = preds[i]
                loss = self.cal_loss(pred, gt_masks)
                loss_total = loss_total + loss
            return loss_total

        else:
            loss = self.cal_loss(preds, gt_masks)
            return loss


    def lossIoU(self, preds, gt_masks):
        if isinstance(preds, list):
            loss_total = 0
            for i in range(len(preds)):
                pred = preds[i]
                gt_mask = gt_masks[i]
                loss = self.softiou(pred, gt_mask)
                loss_total = loss_total + loss
            return loss_total / len(preds)

        elif isinstance(preds, tuple):
            a = []
            for i in range(len(preds)):
                pred = preds[i]
                loss = self.softiou(pred, gt_masks)
                a.append(loss)
            loss_total = a[0] + a[1] + a[2] + a[3] + a[4] + a[5]
            return loss_total

        else:
            loss = self.softiou(preds, gt_masks)
            return loss




    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        if isinstance(m, nn.Conv2d):
            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
            fan_out //= m.groups
            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
            if m.bias is not None:
                m.bias.data.zero_()
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)


if __name__ == '__main__':
    # config_vit = get_SCTrans_config()
    # # model = SCTransNet(config_vit, mode='train', deepsuper=True)
    # model = model

    net = IRSTDNet(model_name='SCTransNet')
    inputs = torch.rand(1, 1, 256, 256)
    output = net(inputs)
    flops, params = profile(net, (inputs,))

    print("-" * 50)
    print('FLOPs = ' + str(flops / 1000 ** 3) + ' G')
    print('Params = ' + str(params / 1000 ** 2) + ' M')