import sys

sys.path.append("./nets")
from feature_extract import ResGCNFeatureExtract, ViTFeatureExtract
from feature_fuse import AttentionFuse, SelfGateV1, SelfGateV2
from global_assemblers import COAssembler, HLMAssembler, KruskalAssembler
from matching import CrossAttnDecoder, MatchingNet
from pairing import PairingNet
from rule_based_alg import align_polygons
from utils.focal_loss import FocalLoss, FocalLoss2
from utils.graph_cluster import graph_cluster
from utils.infornce_loss import InfoNCE

from nets.score_evaluators import (CNNScoreEvaluator, LinearClassify,
                                   MaxDiagScoreEvalutor)

feature_extract_nets = {
    "ResGCN": ResGCNFeatureExtract,
    "ViT": ViTFeatureExtract
}

feature_fuse_nets = {
    "Attention": AttentionFuse,
    "SelfGateV1": SelfGateV1,
    "SelfGateV2": SelfGateV2
}

decoder_nets = {
    "CrossAttn": CrossAttnDecoder
}

score_evaluator = {
    "MaxDiag": MaxDiagScoreEvalutor,
    "CNN": CNNScoreEvaluator
}

global_assemblers = {
    "CO": COAssembler,
    "Kruskal": KruskalAssembler,
    "HLM": HLMAssembler
}
