from typing import Optional
import pandas as pd

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


logger = get_logger(__name__)


def get_score(judgement: dict, criterion: str) -> Optional[int]:
    if judgement and criterion in judgement:
        return judgement[criterion].get("score", None)
    return None

def apply_filtering_to_detail(
    unwrapped_detail_df: pd.DataFrame,
) -> pd.DataFrame:
    logger.info(
        f"Applying filtering to detail dataset with {unwrapped_detail_df['id'].nunique()} samples."
    )

    detail_accuracy = unwrapped_detail_df["judgement"].apply(lambda x: get_score(x, "accuracy"))
    detail_fluency = unwrapped_detail_df["judgement"].apply(lambda x: get_score(x, "fluency"))
    detail_grammar = unwrapped_detail_df["judgement"].apply(lambda x: get_score(x, "grammar"))
    detail_style = unwrapped_detail_df["judgement"].apply(lambda x: get_score(x, "style"))
    detail_completeness = unwrapped_detail_df["judgement"].apply(
        lambda x: get_score(x, "completeness")
    )

    detail_en_wc = unwrapped_detail_df["en"].apply(
        lambda x: word_count(x) if isinstance(x, str) else 0
    )
    detail_pl_wc = unwrapped_detail_df["pl"].apply(
        lambda x: word_count(x) if isinstance(x, str) else 0
    )
    detail_pl_en_wcr = detail_pl_wc / detail_en_wc.replace(0, 1)

    improper_detail_mask = (
        (detail_accuracy < 3)
        | (detail_fluency < 3)
        | (detail_grammar < 3)
        | (detail_style < 3)
        | (detail_completeness < 3)
        | (detail_pl_en_wcr > 1.2)
    )
    improper_detail_ids = unwrapped_detail_df[improper_detail_mask]["id"].unique()
    filtered_detail_df = unwrapped_detail_df[
        ~unwrapped_detail_df["id"].isin(improper_detail_ids)
    ].reset_index(drop=True)

    logger.info(f"Filtered out {len(improper_detail_ids)} improper detail samples.")
    return filtered_detail_df


def apply_filtering_to_complex_reasoning(
    unwrapped_complex_reasoning_df: pd.DataFrame,
) -> pd.DataFrame:
    logger.info(
        f"Applying filtering to complex reasoning dataset with {unwrapped_complex_reasoning_df['id'].nunique()} samples."
    )

    complex_reasoning_accuracy = unwrapped_complex_reasoning_df["judgement"].apply(
        lambda x: get_score(x, "accuracy")
    )
    complex_reasoning_fluency = unwrapped_complex_reasoning_df["judgement"].apply(
        lambda x: get_score(x, "fluency")
    )
    complex_reasoning_grammar = unwrapped_complex_reasoning_df["judgement"].apply(
        lambda x: get_score(x, "grammar")
    )
    complex_reasoning_style = unwrapped_complex_reasoning_df["judgement"].apply(
        lambda x: get_score(x, "style")
    )
    complex_reasoning_completeness = unwrapped_complex_reasoning_df["judgement"].apply(
        lambda x: get_score(x, "completeness")
    )

    improper_complex_reasoning_mask = (
        (complex_reasoning_accuracy < 3)
        | (complex_reasoning_fluency < 3)
        | (complex_reasoning_grammar < 3)
        | (complex_reasoning_style < 3)
        | (complex_reasoning_completeness < 3)
    )
    improper_complex_reasoning_ids = unwrapped_complex_reasoning_df[
        improper_complex_reasoning_mask
    ]["id"].unique()
    filtered_complex_reasoning_df = unwrapped_complex_reasoning_df[
        ~unwrapped_complex_reasoning_df["id"].isin(improper_complex_reasoning_ids)
    ].reset_index(drop=True)

    logger.info(
        f"Filtered out {len(improper_complex_reasoning_ids)} improper complex reasoning samples."
    )
    return filtered_complex_reasoning_df


def apply_filtering_to_conversation(
    unwrapped_conversation_df: pd.DataFrame,
) -> pd.DataFrame:
    logger.info(
        f"Applying filtering to conversation dataset with {unwrapped_conversation_df['id'].nunique()} samples."
    )

    conversation_accuracy = unwrapped_conversation_df["judgement"].apply(
        lambda x: get_score(x, "accuracy")
    )
    conversation_fluency = unwrapped_conversation_df["judgement"].apply(
        lambda x: get_score(x, "fluency")
    )
    conversation_grammar = unwrapped_conversation_df["judgement"].apply(
        lambda x: get_score(x, "grammar")
    )
    conversation_style = unwrapped_conversation_df["judgement"].apply(
        lambda x: get_score(x, "style")
    )
    conversation_completeness = unwrapped_conversation_df["judgement"].apply(
        lambda x: get_score(x, "completeness")
    )

    improper_conversation_mask = (
        (conversation_accuracy < 3)
        | (conversation_fluency < 3)
        | (conversation_grammar < 3)
        | (conversation_style < 3)
        | (conversation_completeness < 3)
    )
    improper_conversation_ids = unwrapped_conversation_df[improper_conversation_mask][
        "id"
    ].unique()
    filtered_conversation_df = unwrapped_conversation_df[
        ~unwrapped_conversation_df["id"].isin(improper_conversation_ids)
    ].reset_index(drop=True)

    logger.info(f"Filtered out {len(improper_conversation_ids)} improper conversation samples.")
    return filtered_conversation_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 os
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--detail_path", type=str, required=True)
    parser.add_argument("--complex_reasoning_path", type=str, required=True)
    parser.add_argument("--conversation_path", type=str, required=True)
    parser.add_argument("--output_detail_path", type=str, required=True)
    parser.add_argument("--output_complex_reasoning_path", type=str, required=True)
    parser.add_argument("--output_conversation_path", type=str, required=True)
    parser.add_argument("--images_source", type=str, default="")
    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)
    args = parser.parse_args()

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

    logger.info(f"Loading complex reasoning data from {args.complex_reasoning_path}")
    complex_reasoning_df = pd.read_json(
        args.complex_reasoning_path,
        lines=args.complex_reasoning_path.endswith(".jsonl"),
    )

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

    unwrapped_detail_df = detail_df if args.unwrapped else unwrap(detail_df)
    unwrapped_complex_reasoning_df = (
        complex_reasoning_df if args.unwrapped else unwrap(complex_reasoning_df)
    )
    unwrapped_conversation_df = conversation_df if args.unwrapped else unwrap(conversation_df)

    processed_unwrapped_detail_df = apply_filtering_to_detail(unwrapped_detail_df)
    processed_unwrapped_complex_reasoning_df = apply_filtering_to_complex_reasoning(
        unwrapped_complex_reasoning_df
    )
    processed_unwrapped_conversation_df = apply_filtering_to_conversation(
        unwrapped_conversation_df
    )

    processed_en_unwrapped_detail_df, processed_pl_unwrapped_detail_df = (
        apply_language_subsampling(
            processed_unwrapped_detail_df, args.en_frac, random_state=args.random_state
        )
    )

    (
        processed_en_unwrapped_complex_reasoning_df,
        processed_pl_unwrapped_complex_reasoning_df,
    ) = apply_language_subsampling(
        processed_unwrapped_complex_reasoning_df,
        args.en_frac,
        random_state=args.random_state,
    )

    processed_en_unwrapped_conversation_df, processed_pl_unwrapped_conversation_df = (
        apply_language_subsampling(
            processed_unwrapped_conversation_df,
            args.en_frac,
            random_state=args.random_state,
        )
    )

    logger.info("Wrapping processed data back to original format")
    processed_en_detail_df = wrap(processed_en_unwrapped_detail_df, lang="en")
    processed_pl_detail_df = wrap(processed_pl_unwrapped_detail_df, lang="pl")

    processed_en_complex_reasoning_df = wrap(processed_en_unwrapped_complex_reasoning_df, lang="en")
    processed_pl_complex_reasoning_df = wrap(processed_pl_unwrapped_complex_reasoning_df, lang="pl")

    processed_en_conversation_df = wrap(processed_en_unwrapped_conversation_df, lang="en")
    processed_pl_conversation_df = wrap(processed_pl_unwrapped_conversation_df, lang="pl")

    processed_detail_df = pd.concat(
        [processed_en_detail_df, processed_pl_detail_df],
        ignore_index=True,
    )
    processed_complex_reasoning_df = pd.concat(
        [
            processed_en_complex_reasoning_df,
            processed_pl_complex_reasoning_df,
        ],
        ignore_index=True,
    )
    processed_conversation_df = pd.concat(
        [
            processed_en_conversation_df,
            processed_pl_conversation_df,
        ],
        ignore_index=True,
    )

    processed_detail_df = apply_images_source(processed_detail_df, args.images_source)
    processed_complex_reasoning_df = apply_images_source(
        processed_complex_reasoning_df, args.images_source
    )
    processed_conversation_df = apply_images_source(processed_conversation_df, args.images_source)

    processed_detail_df = apply_string_id(processed_detail_df)
    processed_complex_reasoning_df = apply_string_id(processed_complex_reasoning_df)
    processed_conversation_df = apply_string_id(processed_conversation_df)

    os.makedirs(os.path.dirname(args.output_detail_path), exist_ok=True)
    os.makedirs(os.path.dirname(args.output_complex_reasoning_path), exist_ok=True)
    os.makedirs(os.path.dirname(args.output_conversation_path), exist_ok=True)

    logger.info(f"Saving processed detail data to {args.output_detail_path}")
    processed_detail_df.to_json(
        args.output_detail_path,
        lines=args.output_detail_path.endswith(".jsonl"),
        orient="records",
    )
    logger.info(
        f"Saving processed complex reasoning data to {args.output_complex_reasoning_path}"
    )
    processed_complex_reasoning_df.to_json(
        args.output_complex_reasoning_path,
        lines=args.output_complex_reasoning_path.endswith(".jsonl"),
        orient="records",
    )
    logger.info(f"Saving processed conversation data to {args.output_conversation_path}")
    processed_conversation_df.to_json(
        args.output_conversation_path,
        lines=args.output_conversation_path.endswith(".jsonl"),
        orient="records",
    )
