import pandas as pd
from comet import download_model, load_from_checkpoint
from src.logger import get_logger
from src.generator import ModelConfig, LLMGenerator


logger = get_logger(__name__)


def translate(df: pd.DataFrame, translator_config: ModelConfig) -> pd.DataFrame:
    translation_prompt = lambda sample: translator_config.prompt_template.format(text=sample)

    translated_df = df.copy()
    with LLMGenerator(translator_config) as translator:
        en_captions = translated_df["caption"].tolist()
        pl_captions = translator.generate(en_captions, prompt=translation_prompt)

    translated_df["caption"] = pl_captions
    return translated_df


def score_comet(base_df: pd.DataFrame, translated_df: pd.DataFrame, model_name: str, num_gpus:int = 1) -> pd.DataFrame:
    en_captions = base_df["caption"].tolist()
    pl_captions = translated_df["caption"].tolist()
    data = [{"src": src, "mt": mt} for src, mt in zip(en_captions, pl_captions)]

    model = load_from_checkpoint(download_model(model_name))
    model_output = model.predict(data, batch_size=32, gpus=num_gpus)
    
    translated_df["comet_score"] = model_output[0]
    return translated_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 JSONL file."
    )
    parser.add_argument("--output_path", type=str, required=True, help="Path to the output file.")
    parser.add_argument(
        "--translator_config",
        type=str,
        default=None,
        help="Path to the model configuration file.",
    )
    parser.add_argument(
        "--comet_model",
        type=str,
        default="Unbabel/wmt23-cometkiwi-da-xl",
        help="Name of the COMET model to use for scoring.",
    )
    args = parser.parse_args()

    logger.info(f"Loading data from {args.input_path}")
    data_df = pd.read_json(args.input_path, lines=True)
    # data_df = data_df.sample(n=50, random_state=42).reset_index(drop=True)

    logger.info(f"Loading translator configuration from {args.translator_config}")
    translator_config = ModelConfig.from_yaml(args.translator_config)

    logger.info("Translating data")
    translated_df = translate(data_df, translator_config)
    translated_df = data_df.copy()

    # logger.info("Scoring translated data with COMET")
    # translated_df = score_comet(data_df, translated_df, args.comet_model, num_gpus=translator_config.num_gpus)

    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, orient="records", lines=use_lines, indent=2 if not use_lines else None
    )
