import pandas as pd
from src.datasets.utils import unwrap, wrap
from src.logger import get_logger
from src.generator import ModelConfig, LLMGenerator
from src.datasets.utils import calculate_comet_score


logger = get_logger(__name__)


def translate(
    df: pd.DataFrame, translator_config: ModelConfig, settings: tuple[bool, bool]
) -> pd.DataFrame:
    unwrapped_df = unwrap(df).copy()

    translation_prompt = lambda sample: translator_config.prompt_template.format(text=sample)

    from_human_values_idx = unwrapped_df["from"].apply(lambda x: x == "human")
    from_gpt_values_idx = unwrapped_df["from"].apply(lambda x: x == "gpt")

    translate_from_human, translate_from_gpt = settings
    from_human_values = unwrapped_df["value"][from_human_values_idx]
    from_gpt_values = unwrapped_df["value"][from_gpt_values_idx]
    

    with LLMGenerator(translator_config) as translator:
        from_human_translations, from_gpt_translations = from_human_values, from_gpt_values

        if translate_from_human:
            logger.info(f"Translating {len(from_human_values)} from human values.")
            from_human_translations = translator.generate(from_human_values.tolist(), prompt=translation_prompt)

        if translate_from_gpt:
            logger.info(f"Translating {len(from_gpt_values)} from gpt values.")
            from_gpt_translations = translator.generate(from_gpt_values.tolist(), prompt=translation_prompt)

        unwrapped_df.loc[from_human_values_idx, "value"] = from_human_translations
        unwrapped_df.loc[from_gpt_values_idx, "value"] = from_gpt_translations

    wrapped_df = wrap(unwrapped_df)
    return wrapped_df


def score_comet(
    base_df: pd.DataFrame, translated_df: pd.DataFrame, settings: tuple[bool, bool]
) -> pd.DataFrame:
    unwrapped_base_df = unwrap(base_df).copy()
    unwrapped_translated_df = unwrap(translated_df).copy()

    from_human_idx = unwrapped_translated_df["from"].apply(lambda x: x == "human")
    from_gpt_idx = unwrapped_translated_df["from"].apply(lambda x: x == "gpt")

    from_human_values = unwrapped_base_df["value"][from_human_idx].tolist()
    from_gpt_values = unwrapped_base_df["value"][from_gpt_idx].tolist()

    from_human_translations = unwrapped_translated_df["value"][from_human_idx].tolist()
    from_gpt_translations = unwrapped_translated_df["value"][from_gpt_idx].tolist()

    translate_from_human, translate_from_gpt = settings

    if translate_from_human:
        logger.info(f"Calculating COMET scores for {len(from_human_values)} from human translations.")
        from_human_scores = calculate_comet_score(from_human_values, from_human_translations)

        unwrapped_translated_df.loc[from_human_idx, "metadata"] = [
            {**(md if isinstance(md, dict) else {}), "comet_score": score} 
            for md, score in zip(unwrapped_translated_df.loc[from_human_idx, "metadata"], from_human_scores)
        ]

    if translate_from_gpt:
        logger.info(f"Calculating COMET scores for {len(from_gpt_values)} from gpt translations.")
        from_gpt_scores = calculate_comet_score(from_gpt_values, from_gpt_translations)

        unwrapped_translated_df.loc[from_gpt_idx, "metadata"] = [
            {**(md if isinstance(md, dict) else {}), "comet_score": score} 
            for md, score in zip(unwrapped_translated_df.loc[from_gpt_idx, "metadata"], from_gpt_scores)
        ]

    wrapped_translated_df = wrap(unwrapped_translated_df)
    return wrapped_translated_df


def judge(
    base_df: pd.DataFrame, translated_df: pd.DataFrame, judge_config: ModelConfig, settings: tuple[bool, bool]
) -> pd.DataFrame:
    unwrapped_base_df = unwrap(base_df).copy()
    unwrapped_translated_df = unwrap(translated_df).copy()

    from_human_idx = unwrapped_translated_df["from"].apply(lambda x: x == "human")
    from_gpt_idx = unwrapped_translated_df["from"].apply(lambda x: x == "gpt")

    from_human_values = unwrapped_base_df["value"][from_human_idx].tolist()
    from_gpt_values = unwrapped_base_df["value"][from_gpt_idx].tolist()

    from_human_translations = unwrapped_translated_df["value"][from_human_idx].tolist()
    from_gpt_translations = unwrapped_translated_df["value"][from_gpt_idx].tolist()

    translate_from_human, translate_from_gpt = settings

    score_prompt = lambda sample: judge_config.prompt_template.format(
        src_lang="English", target_lang="Polish", src_text=sample[0], target_text=sample[1]
    )

    with LLMGenerator(judge_config) as judge:
        if translate_from_human:
            logger.info(f"Scoring {len(from_human_values)} from human translations.")
            from_human_scores = judge.generate(
                list(zip(from_human_values, from_human_translations)),
                prompt=score_prompt,
            )
            unwrapped_translated_df.loc[from_human_idx, "metadata"] = [
                {**(md if isinstance(md, dict) else {}), "score": score} 
                for md, score in zip(unwrapped_translated_df.loc[from_human_idx, "metadata"], from_human_scores)
            ]

        if translate_from_gpt:
            logger.info(f"Scoring {len(from_gpt_values)} from gpt translations.")
            from_gpt_scores = judge.generate(
                list(zip(from_gpt_values, from_gpt_translations)),
                prompt=score_prompt,
            )
            unwrapped_translated_df.loc[from_gpt_idx, "metadata"] = [
                {**(md if isinstance(md, dict) else {}), "score": score}
                for md, score in zip(unwrapped_translated_df.loc[from_gpt_idx, "metadata"], from_gpt_scores)
            ]

    wrapped_translated_df = wrap(unwrapped_translated_df)
    return wrapped_translated_df


def filter_comet(
    translated_df: pd.DataFrame, filter_settings: tuple[float, float]
) -> pd.DataFrame:
    from_human_comet_threshold, from_gpt_comet_threshold = filter_settings
    logger.info(f"From human COMET threshold: {from_human_comet_threshold}")
    logger.info(f"From gpt COMET threshold: {from_gpt_comet_threshold}")

    all_filtered_conversations = []
    all_qa_pairs, filtered_qa_pairs = 0, 0

    for _, row in translated_df.iterrows():
        conversations = row.get("conversations", [])
        filtered_conversations = []
        for i in range(0, len(conversations), 2):
            all_qa_pairs += 2

            human_message = conversations[i]
            gpt_message = conversations[i + 1]

            if not human_message.get("from", "") == "human" or not gpt_message.get("from", "") == "gpt":
                logger.warning(f"Unexpected conversation format at row {_}")

            human_comet_score = conversations[i].get("metadata", {}).get("comet_score", None)
            gpt_comet_score = conversations[i + 1].get("metadata", {}).get("comet_score", None)
            
            if (
                (human_comet_score and human_comet_score < from_human_comet_threshold) or 
                (gpt_comet_score and gpt_comet_score < from_gpt_comet_threshold)
            ):
                continue

            filtered_qa_pairs += 2

            filtered_conversations.append(conversations[i])
            filtered_conversations.append(conversations[i + 1])

        if filtered_conversations and not any([m for m in filtered_conversations if m['from'] == 'human' and '<image>' in m['value']]):
            # sometimes, we can filter out the first QA pair which contains the image reference, thus we need to add it back
            filtered_conversations[0]['value'] = '<image>\n' + filtered_conversations[0]['value']

        all_filtered_conversations.append(filtered_conversations)

    filtered_df = translated_df.copy()
    filtered_df["conversations"] = all_filtered_conversations

    # get rid of rows with empty conversations
    filtered_df = filtered_df[filtered_df["conversations"].map(len) > 0].reset_index(drop=True)

    logger.info(f"Filtered out {all_qa_pairs - filtered_qa_pairs} out of {all_qa_pairs} QA pairs.")

    return filtered_df
                


if __name__ == "__main__":
    import os
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--input_path", type=str, required=True, help="Path to the input file.")
    parser.add_argument("--output_path", type=str, required=True, help="Path to the output file.")
    parser.add_argument("--skip_translate", action="store_true", help="Whether to translate the data.")
    parser.add_argument("--skip_from_human", action="store_true", help="Whether to translate the question.")
    parser.add_argument("--skip_from_gpt", action="store_true", help="Whether to translate the response.")
    parser.add_argument("--translator_config", type=str, default=None, help="Path to the model configuration file.")
    parser.add_argument("--comet", action="store_true", help="Whether to use COMET scoring.")
    parser.add_argument("--judge", action="store_true", help="Whether to use a judge model.")
    parser.add_argument("--judge_config", type=str)
    parser.add_argument("--filter", action="store_true", help="Whether to filter out bad translations.")
    parser.add_argument("--from_human_comet_threshold", type=float, default=0.0, help="Threshold for from human scores.")
    parser.add_argument("--from_gpt_comet_threshold", type=float, default=0.0, help="Threshold for from gpt scores.")
    args = parser.parse_args()

    logger.info(f"Loading data from {args.input_path}")
    df = pd.read_json(args.input_path, lines=args.input_path.endswith(".jsonl"))

    translated_df = df
    settings = (not args.skip_from_human, not args.skip_from_gpt)

    if not args.skip_translate:
        logger.info("Translating data.")
        logger.info(f"Loading translator model configuration from {args.translator_config}")
        translator_config = ModelConfig.from_yaml(args.translator_config)
        translated_df = translate(df, translator_config, settings)

    if args.comet:
        logger.info("Scoring translations with COMET.")
        translated_df = score_comet(df, translated_df, settings)

    if args.judge:
        logger.info("Judging translations.")
        logger.info("Using judge model configuration same as translator.")
        judge_config = ModelConfig.from_yaml(args.judge_config)
        translated_df = judge(df, translated_df, judge_config, settings)

    if args.filter:
        logger.info("Filtering out bad translations by COMET scores.")
        filter_settings = (args.from_human_comet_threshold, args.from_gpt_comet_threshold)
        translated_df = filter_comet(translated_df, filter_settings)

    logger.info(f"Saving translated data to {args.output_path}")
    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
    use_lines = args.output_path.endswith(".jsonl")
    translated_df.to_json(
        args.output_path, lines=use_lines, orient="records", indent=2 if not use_lines else None
    )
