import pandas as pd
from src.logger import get_logger

logger = get_logger(__name__)

def subsample(first_df: pd.DataFrame, second_df: pd.DataFrame, second_ratio: float = 0.15, seed: int = 42) -> pd.DataFrame:
    if second_ratio <= 0 or second_ratio >= 1:
        raise ValueError("second_ratio must be between 0 and 1")
    
    logger.info(f"First dataset size: {len(first_df)}")
    logger.info(f"Second dataset size: {len(second_df)}")

    second_indices = second_df["id"].sample(frac=second_ratio, random_state=seed)
    first_indices = first_df[~first_df["id"].isin(second_indices)]["id"]

    logger.info(f"Selected {len(second_indices)} examples from the second dataset")
    logger.info(f"Selected {len(first_indices)} examples from the first dataset")

    combined_df = pd.concat([first_df[first_df["id"].isin(first_indices)], second_df[second_df["id"].isin(second_indices)]])
    combined_df = combined_df.sample(frac=1, random_state=seed).reset_index(drop=True)

    final_cols = first_df.columns
    return combined_df[final_cols]

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

    parser = argparse.ArgumentParser()
    parser.add_argument("--first_path", type=str, required=True)
    parser.add_argument("--second_path", type=str, required=True)
    parser.add_argument("--output_path", type=str, required=True)
    parser.add_argument("--second_ratio", type=float, default=0.15)
    parser.add_argument("--seed", type=int, default=42)

    args = parser.parse_args()

    logger.info(f"Loading first dataset from {args.first_path}")
    logger.info(f"Loading second dataset from {args.second_path}")

    first_df = pd.read_json(args.first_path, lines=args.first_path.endswith(".jsonl"))
    second_df = pd.read_json(args.second_path, lines=args.second_path.endswith(".jsonl"))

    combined_df = subsample(first_df, second_df, args.second_ratio, args.seed)

    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
    combined_df.to_json(args.output_path, orient="records", lines=args.output_path.endswith(".jsonl"))
    logger.info(f"Combined dataset written to {args.output_path}")
    