import pandas as pd
from src.datasets.wit.utils import get_subcategories
from src.logger import get_logger

logger = get_logger(__name__)


def train_test_split(df: pd.DataFrame, categories_df: pd.DataFrame, test_size: float, seed: int) -> tuple[pd.DataFrame, pd.DataFrame]:
    logger.info(f"Starting train-test split with test size {test_size} and seed {seed}")
    _, polish_pages = get_subcategories(categories_df, 'Kategoria:Polska')

    polish_subcategory_df = df[df['page_title'].isin(polish_pages)]
    non_polish_subcategory_df = df[~df['page_title'].isin(polish_pages)]

    logger.info(f"Polish subcategory size: {len(polish_subcategory_df)}")
    logger.info(f"Non-Polish subcategory size: {len(non_polish_subcategory_df)}")

    test_polish_df = pd.concat([
        polish_subcategory_df.sample(frac=test_size, random_state=seed),
        non_polish_subcategory_df.sample(frac=test_size, random_state=seed)
    ]).reset_index(drop=True)

    train_polish_df = df.drop(test_polish_df.index).reset_index(drop=True)
    logger.info(f"Train size: {len(train_polish_df)}, Test size: {len(test_polish_df)}")

    return train_polish_df, test_polish_df


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input_path", type=str, required=True, help="Path to the processed WIT data file"
    )
    parser.add_argument(
        "--categories_path",
        type=str,
        required=True,
        help="Path to the file containing list of categories",
    )
    parser.add_argument(
        "--output_train_path",
        type=str,
        required=True,
        help="Path to save the training split of the WIT data",
    )
    parser.add_argument(
        "--output_test_path",
        type=str,
        required=True,
        help="Path to save the test split of the WIT data",
    )
    parser.add_argument(
        "--test_size",
        type=float,
        default=0.1,
        help="Proportion of the dataset to include in the test split",
    )
    parser.add_argument(
        "--seed", type=int, default=42, help="Random seed for reproducibility"
    )
    args = parser.parse_args()

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

    logger.info(f"Loading category tree from {args.categories_path}")
    category_tree_df = pd.read_json(args.categories_path, lines=args.categories_path.endswith(".jsonl"))

    train_df, test_df = train_test_split(df, category_tree_df, args.test_size, args.seed)

    logger.info(f"Saving training split to {args.output_train_path}")
    use_lines = args.output_train_path.endswith(".jsonl")
    train_df.to_json(args.output_train_path, lines=use_lines, orient="records", indent=2 if not use_lines else None)

    logger.info(f"Saving test split to {args.output_test_path}")
    use_lines = args.output_test_path.endswith(".jsonl")
    test_df.to_json(args.output_test_path, lines=use_lines, orient="records", indent=2 if not use_lines else None)