import os
import logging
import torch.distributed as dist
import torch
import pandas as pd
import random
from transformers.models.siglip.tokenization_siglip import SiglipTokenizer
from transformers.models.gemma.tokenization_gemma_fast import GemmaTokenizerFast
from transformers.models.gemma.tokenization_gemma import GemmaTokenizer

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
logging.basicConfig(
    format="[%(levelname)s|%(filename)s:%(lineno)s] %(asctime)s >> %(message)s"
)


def is_main_process():
    if dist.is_initialized():
        return dist.get_rank() == 0
    else:
        return True
    

def get_local_rank():
    if dist.is_initialized():
        return int(os.getenv('LOCAL_RANK'))
    else:
        return 0


def get_logger(name):
    logger = logging.getLogger(name)
    logger.setLevel(logging.INFO)
    return logger


def is_debug():
    return int(os.environ.get('DEBUG', 0))


def get_mbeir_instruction_row(sample_row, inst_df):
    dataset_id = int(sample_row.data_id.split(':')[0])
    qmod = []
    if not pd.isna(sample_row.image_path):
        qmod.append('image')
    if not pd.isna(sample_row.question):
        qmod.append('text')
    qmod = ','.join(qmod)

    pmod = []
    assert isinstance(sample_row.passage_id, list)
    if not pd.isna(sample_row.passage_image_path).all():
        pmod.append('image')
    assert isinstance(sample_row.passage_text, list)
    if not pd.isna(sample_row.passage_text).all():
        pmod.append('text')
    pmod = ','.join(pmod)

    inst_row = inst_df[(inst_df.dataset_id == dataset_id) & (
        inst_df.query_modality == qmod) & (inst_df.cand_modality == pmod)]
    assert len(inst_row) == 1  
    return inst_row


def sample_mbeir_instruction(row, inst_df):
    inst_row = get_mbeir_instruction_row(row, inst_df).iloc[0]
    inst = random.choice(
        [inst_row.prompt_1, inst_row.prompt_2, inst_row.prompt_3, inst_row.prompt_4])
    return inst


def rank_collapse(x: torch.Tensor):
    with torch.no_grad():
        dim = 1 if len(x.shape) == 3 else 0
        mu = torch.linalg.matrix_norm(x - x.mean(dim=dim, keepdim=True), ord='fro')
        return mu / torch.linalg.matrix_norm(x, ord='fro')


def is_siglip_tokenizer(tokenizer):
    return isinstance(tokenizer, (SiglipTokenizer, GemmaTokenizer, GemmaTokenizerFast))


def is_siglip2_tokenizer(tokenizer):
    return isinstance(tokenizer, (GemmaTokenizer, GemmaTokenizerFast))
