## TODO: 实现一个映射层，将输出得到的特征维度映射到embed空间当中
from models.proj_models import LLM2EmbedProjector, LLM2EmbedAddProjector
from typing import Optional
from logger import get_logger
from transformers.modeling_utils import load_state_dict
from transformers.trainer_utils import get_last_checkpoint
import os
from . import register_loader
logger = get_logger(__name__)


@register_loader("load_mlp_projector")
def load_mlp_projector(
        fc_args: "FCArguments",
        in_dim = 3074, 
        is_trainable: Optional[bool] = False,
        resume_from_checkpoint=None,
        output_dir=None
):
    if fc_args.stage == "classify":
        mlp_model = LLM2EmbedProjector(
            in_dim = in_dim,
            out_dim = fc_args.embed_dim, 
            fp16=False
        )

        if resume_from_checkpoint:
            if os.path.exists(output_dir):
                new_model_name_or_path_proj = get_last_checkpoint(output_dir)
                logger.info("start to load pfc not from: {}, but from {}".format(
                    fc_args.model_name_or_path_proj, new_model_name_or_path_proj
                ))
            else:
                new_model_name_or_path_proj = None
                logger.info('not exists (output_dir) for pfc {}, load pfc from {}'.format(output_dir, fc_args.model_name_or_path_proj))

            if new_model_name_or_path_proj:
                logger.info("(resume process) start to load pfc from: {}".format(new_model_name_or_path_proj))
                mlp_model.maybe_load_state_dict_from_path(new_model_name_or_path_proj)
            else:
                logger.info("(resume failed) still start to load pfc from: {}".format(fc_args.model_name_or_path_proj))
                mlp_model.maybe_load_state_dict_from_path(fc_args.model_name_or_path_proj)
        else:
            logger.info("start to load pfc from: {}".format(fc_args.model_name_or_path_proj))
            mlp_model.maybe_load_state_dict_from_path(fc_args.model_name_or_path_proj)


        if is_trainable:
            mlp_model.train()
        else:
            mlp_model.eval()
    else:
        mlp_model = None
    return mlp_model


@register_loader("load_add_projector")
def load_add_projector(
        fc_args: "FCArguments",
        in_dim = 3074, 
        is_trainable: Optional[bool] = False,
        resume_from_checkpoint=None,
        output_dir=None
):
    if fc_args.stage == "classify":
        mlp_model = LLM2EmbedAddProjector(
            in_dim = in_dim,
            out_dim = fc_args.embed_dim, 
            fp16=False
        )
        logger.info("No weights for LLM2EmbedAddProjector")
        if is_trainable:
            mlp_model.train()
        else:
            mlp_model.eval()
    else:
        mlp_model = None
    return mlp_model