import torch
import torch.nn as nn
import torch.nn.functional as F

from   utils import CONFIG
from   networks import m2ms, ops
import sys
sys.path.insert(0, './segment-anything')
from segment_anything import sam_model_registry

class sam_m2m(nn.Module):
    def __init__(self, m2m):
        super(sam_m2m, self).__init__()
        if m2m not in m2ms.__all__:
            raise NotImplementedError("Unknown M2M {}".format(m2m))
        self.m2m = m2ms.__dict__[m2m](nc=256)
        self.seg_model = sam_model_registry['vit_b'](checkpoint=None)
        self.seg_model.eval()

    def forward(self, image, guidance):
        self.seg_model.eval()
        with torch.no_grad():
            feas, masks = self.seg_model.forward_m2m(image, guidance, multimask_output=True)
        pred = self.m2m(feas, image, masks)
        return pred

    def forward_inference(self, image_dict):
        self.seg_model.eval()
        with torch.no_grad():
            feas, masks, post_masks = self.seg_model.forward_m2m_inference(image_dict, multimask_output=True)
        pred = self.m2m(feas, image_dict["image"], masks)
        return feas, pred, post_masks

def get_generator_m2m(seg, m2m):
    if seg == 'sam':
        generator = sam_m2m(m2m=m2m)
    return generator