from typing import Optional
import pandas as pd

from src.datasets.images import ImagesDataset
from src.datasets.llava.utils import unwrap, wrap, word_count
from src.logger import get_logger


logger = get_logger(__name__)


def apply_gpt_translations_filtering(
    unwrapped_df: pd.DataFrame,
) -> pd.DataFrame:
    logger.info(f"Starting filtering on {unwrapped_df['id'].nunique()} samples")

    from_gpt = unwrapped_df["from"] == "gpt"

    pl_wc = unwrapped_df["pl"].apply(
        lambda x: word_count(x) if isinstance(x, str) else 0
    )
    pl_cc = unwrapped_df["pl"].apply(lambda x: len(x) if isinstance(x, str) else 0)
    en_wc = unwrapped_df["en"].apply(word_count)
    en_pl_wcr = en_wc / pl_wc.replace(0, 1)

    improper_translations_mask = (
        (pl_wc == 0)
        | (pl_wc > 42)
        | ((en_wc <= 10) & (en_pl_wcr < 0.2))
        | ((en_wc > 10) & (en_pl_wcr < 0.5))
        | (en_pl_wcr > 2.5)
        | (pl_cc > 200)
    )

    improper_ids = unwrapped_df[improper_translations_mask & from_gpt]["id"].unique()
    logger.info(f"Filtering out {len(improper_ids)} improper samples")

    filtered_df = unwrapped_df[~unwrapped_df["id"].isin(improper_ids)].reset_index(
        drop=True
    )
    return filtered_df

def apply_human_manual_translations(unwrapped_df: pd.DataFrame) -> pd.DataFrame:
    logger.info(f"Applying manual translations on {unwrapped_df['id'].nunique()} samples")

    instructions_en_pl = {
        "Render a clear and concise summary of the photo." : "Przedstaw jasne i zwięzłe podsumowanie zdjęcia.",
        "Share a concise interpretation of the image provided." : "Podaj zwięzłą interpretację załączonego obrazu.",
        "Write a terse but informative summary of the picture.": "Napisz zwięzłe, ale informacyjne podsumowanie zdjęcia.",
        "Give a brief description of the image.": "Podaj krótki opis obrazu.",
        "Provide a brief description of the given image.": "Podaj krótki opis danego obrazu.",
        "What is in the photo?": "Co znajduje się na zdjęciu?",
        "What is this?": "Co to jest?",
        "Summarize the visual content of the image.": "Podsumuj treść wizualną obrazu.",
        "Describe the image concisely.": "Opisz zwięźle obraz.",
        "Present a compact description of the photo's key features.": "Przedstaw zwięzły opis kluczowych cech zdjęcia.",
        "Give a short and clear explanation of the subsequent image.": "Podaj krótkie i jasne wyjaśnienie kolejnego obrazu.",
    }

    from_human = unwrapped_df["from"] == "human"
    unwrapped_df.loc[from_human, "pl"] = unwrapped_df.loc[from_human, "en"].map(instructions_en_pl)

    # assert all manual translations are applied
    missing_manual_translations = unwrapped_df[from_human & unwrapped_df["pl"].isna()]["id"].nunique()
    if missing_manual_translations > 0:
        raise ValueError(f"Missing manual translations for {missing_manual_translations} samples")

    return unwrapped_df

def apply_language_subsampling(
    unwrapped_df: pd.DataFrame, en_frac: float, random_state: int = 42
) -> tuple[pd.DataFrame, pd.DataFrame]:
    all_ids = unwrapped_df["id"].drop_duplicates()

    en_ids = all_ids.sample(frac=en_frac, random_state=random_state)
    pl_ids = all_ids[~all_ids.isin(en_ids)]

    logger.info(
        f"Subsampling to {len(en_ids)} English and {len(pl_ids)} Polish samples"
    )
    en_unwrapped_df = unwrapped_df[unwrapped_df["id"].isin(en_ids)].copy()
    pl_unwrapped_df = unwrapped_df[unwrapped_df["id"].isin(pl_ids)].copy()

    return en_unwrapped_df, pl_unwrapped_df


def apply_images_source(df: pd.DataFrame, source: Optional[str] = None) -> pd.DataFrame:
    logger.info(f"Applying images source: {source}")
    df["image_source"] = source
    return df

def apply_string_id(df: pd.DataFrame) -> pd.DataFrame:
    logger.info("Converting ids to strings")
    df["id"] = df["id"].apply(lambda x: f"llava_{x}")
    return df


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--input_path", type=str, required=True)
    parser.add_argument("--output_path", type=str, required=True)
    parser.add_argument("--en_frac", type=float, default=0.15)
    parser.add_argument("--unwrapped", action="store_true")
    parser.add_argument("--random_state", type=int, default=42)
    parser.add_argument("--images_source", type=str, default="llava-pretrain-images")
    args = parser.parse_args()

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

    unwrapped_pretrain_df = unwrap(pretrain_df) if not args.unwrapped else pretrain_df

    processed_unwrapped_pretrain_df = apply_gpt_translations_filtering(unwrapped_pretrain_df)
    processed_unwrapped_pretrain_df = apply_human_manual_translations(processed_unwrapped_pretrain_df)

    processed_en_unwrapped_pretrain_df, processed_pl_unwrapped_pretrain_df = (
        apply_language_subsampling(
            processed_unwrapped_pretrain_df, 
            args.en_frac, args.random_state
        )
    )

    processed_en_pretrain_df = wrap(processed_en_unwrapped_pretrain_df, lang='en')
    processed_pl_pretrain_df = wrap(processed_pl_unwrapped_pretrain_df, lang='pl')
    processed_pretrain_df = pd.concat(
        [processed_en_pretrain_df, processed_pl_pretrain_df], ignore_index=True
    )

    processed_pretrain_df = apply_images_source(
        processed_pretrain_df, args.images_source
    )

    processed_pretrain_df = apply_string_id(processed_pretrain_df)

    logger.info(f"Saving processed pretrain data to {args.output_path}")
    processed_pretrain_df.to_json(
        args.output_path,
        orient="records",
        lines=args.output_path.endswith(".jsonl"),
    )
