import torch
from fvcore.nn import FlopCountAnalysis
from einops.einops import rearrange

from src import get_model_cfg
from src.models.backbone import FPN as topicfm_featnet
from src.models.modules import TopicFormer
from src.utils.dataset import read_scannet_gray

from third_party.loftr.src.loftr.utils.cvpr_ds_config import default_cfg
from third_party.loftr.src.loftr.backbone import ResNetFPN_8_2 as loftr_featnet
from third_party.loftr.src.loftr.loftr_module import LocalFeatureTransformer


def feat_net_flops(feat_net, config, input):
    model = feat_net(config)
    model.eval()
    flops = FlopCountAnalysis(model, input)
    feat_c, _ = model(input)
    return feat_c, flops.total() / 1e9


def coarse_model_flops(coarse_model, config, inputs):
    model = coarse_model(config)
    model.eval()
    flops = FlopCountAnalysis(model, inputs)
    return flops.total() / 1e9


if __name__ == '__main__':
    path_img0 = "assets/scannet_sample_images/scene0711_00_frame-001680.jpg"
    path_img1 = "assets/scannet_sample_images/scene0711_00_frame-001995.jpg"
    img0, img1 = read_scannet_gray(path_img0), read_scannet_gray(path_img1)
    img0, img1 = img0.unsqueeze(0), img1.unsqueeze(0)

    # LoFTR
    loftr_conf = dict(default_cfg)
    feat_c0, loftr_featnet_flops0 = feat_net_flops(loftr_featnet, loftr_conf["resnetfpn"], img0)
    feat_c1, loftr_featnet_flops1 = feat_net_flops(loftr_featnet, loftr_conf["resnetfpn"], img1)
    print("FLOPs of feature extraction in LoFTR: {} GFLOPs".format((loftr_featnet_flops0 + loftr_featnet_flops1)/2))
    feat_c0 = rearrange(feat_c0, 'n c h w -> n (h w) c')
    feat_c1 = rearrange(feat_c1, 'n c h w -> n (h w) c')
    loftr_coarse_model_flops = coarse_model_flops(LocalFeatureTransformer, loftr_conf["coarse"], (feat_c0, feat_c1))
    print("FLOPs of coarse matching model in LoFTR: {} GFLOPs".format(loftr_coarse_model_flops))

    # TopicFM
    topicfm_conf = get_model_cfg()
    feat_c0, topicfm_featnet_flops0 = feat_net_flops(topicfm_featnet, topicfm_conf["fpn"], img0)
    feat_c1, topicfm_featnet_flops1 = feat_net_flops(topicfm_featnet, topicfm_conf["fpn"], img1)
    print("FLOPs of feature extraction in TopicFM: {} GFLOPs".format((topicfm_featnet_flops0 + topicfm_featnet_flops1) / 2))
    feat_c0 = rearrange(feat_c0, 'n c h w -> n (h w) c')
    feat_c1 = rearrange(feat_c1, 'n c h w -> n (h w) c')
    topicfm_coarse_model_flops = coarse_model_flops(TopicFormer, topicfm_conf["coarse"], (feat_c0, feat_c1))
    print("FLOPs of coarse matching model in TopicFM: {} GFLOPs".format(topicfm_coarse_model_flops))