import pandas as pd
import numpy as np

# This was used to split the HAP-E data into chunks that would work with
# the MP. The splits are pretty arbitrary and were pretty much just tuned
# so the accuracy values looked sane. The LLM train and test data actually
# comes from different models

def split_ordered_dataset(
    input_path, 
    output_train_path, 
    output_test_path, 
    train_ratio=0.8
):
    """
    Split dataset into train and test sets without shuffling.
    Takes first train_ratio (e.g., 80%) of each category as train,
    and the remaining as test.
    
    Parameters:
    -----------
    input_path : str
        Path to input parquet file
    output_train_path : str
        Path to save train set parquet file
    output_test_path : str
        Path to save test set parquet file
    train_ratio : float, default=0.8
        Proportion of each category to use for training
    """
    
    # Load the dataset
    print("Loading dataset...")
    data = pd.read_parquet(input_path)
    
    # Extract category from doc_id
    print("Extracting categories...")
    data['category'] = data['doc_id'].str.split('_').str[0]
    
    # Define expected categories
    expected_categories = ['acad', 'blog', 'news', 'fic', 'spok', 'tvm']
    
    # Initialize lists to hold train and test dataframes
    train_dfs = []
    test_dfs = []
    
    print("Splitting each category without shuffling...")
    for category in expected_categories:
        # Get data for current category (preserving original order)
        category_data = data[data['category'] == category].copy()
        
        if len(category_data) == 0:
            print(f"Warning: No data for category '{category}'")
            continue
        
        # Calculate split point
        split_idx = int(len(category_data) * train_ratio)
        
        if split_idx == 0 or split_idx == len(category_data):
            print(f"Warning: Category '{category}' has too few samples ({len(category_data)}). "
                  f"Using all for {'train' if split_idx > 0 else 'test'}.")
            if split_idx > 0:
                train_dfs.append(category_data)
            else:
                test_dfs.append(category_data)
            continue
        
        # Split without shuffling: first train_ratio for train, rest for test
        train_cat = category_data.iloc[:split_idx].copy()
        test_cat = category_data.iloc[split_idx:].copy()
        
        print(f"  {category}: {len(train_cat)} train, {len(test_cat)} test "
              f"({len(test_cat)/len(category_data):.1%} test)")
        
        train_dfs.append(train_cat)
        test_dfs.append(test_cat)
    
    # Combine all categories
    print("\nCombining results...")
    
    if train_dfs:
        train_data = pd.concat(train_dfs, ignore_index=False)
    else:
        train_data = pd.DataFrame(columns=data.columns)
        
    if test_dfs:
        test_data = pd.concat(test_dfs, ignore_index=False)
    else:
        test_data = pd.DataFrame(columns=data.columns)
    
    # Remove the temporary category column
    if 'category' in train_data.columns:
        train_data = train_data.drop(columns=['category'])
    if 'category' in test_data.columns:
        test_data = test_data.drop(columns=['category'])
    
    # Save the split datasets
    print("Saving results...")
    train_data.to_parquet(output_train_path, index=False)
    test_data.to_parquet(output_test_path, index=False)
    
    # Print summary statistics
    print("\n" + "="*50)
    print("SPLIT SUMMARY")
    print("="*50)
    print(f"Total samples: {len(data):,}")
    print(f"Train samples: {len(train_data):,} ({len(train_data)/len(data):.1%})")
    print(f"Test samples:  {len(test_data):,} ({len(test_data)/len(data):.1%})")
    
    # Show category distribution in each split
    print("\nCategory distribution:")
    train_data_temp = train_data.copy()
    test_data_temp = test_data.copy()
    train_data_temp['category'] = train_data_temp['doc_id'].str.split('_').str[0]
    test_data_temp['category'] = test_data_temp['doc_id'].str.split('_').str[0]
    
    for category in expected_categories:
        train_count = (train_data_temp['category'] == category).sum()
        test_count = (test_data_temp['category'] == category).sum()
        total = train_count + test_count
        
        if total > 0:
            train_pct = train_count / total * 100 if total > 0 else 0
            test_pct = test_count / total * 100 if total > 0 else 0
            print(f"  {category:5s}: {train_count:5d} train ({train_pct:5.1f}%), "
                  f"{test_count:5d} test ({test_pct:5.1f}%)")
    
    print(f"\nTrain set saved to: {output_train_path}")
    print(f"Test set saved to: {output_test_path}")
    
    return train_data, test_data

# Alternative version that sorts within each category before splitting
def split_ordered_sorted_dataset(
    input_path, 
    output_train_path, 
    output_test_path, 
    train_ratio=0.8,
    sort_by='doc_id'
):
    """
    Split dataset into train and test sets without shuffling.
    Sorts each category by specified column before taking first train_ratio as train.
    
    Parameters:
    -----------
    input_path : str
        Path to input parquet file
    output_train_path : str
        Path to save train set parquet file
    output_test_path : str
        Path to save test set parquet file
    train_ratio : float, default=0.8
        Proportion of each category to use for training
    sort_by : str, default='doc_id'
        Column to sort by within each category before splitting
    """
    
    # Load the dataset
    print("Loading dataset...")
    data = pd.read_parquet(input_path)
    
    # Extract category from doc_id
    print("Extracting categories...")
    data['category'] = data['doc_id'].str.split('_').str[0]
    
    # Define expected categories
    expected_categories = ['acad', 'blog', 'news', 'fic', 'spok', 'tvm']
    
    # Initialize lists to hold train and test dataframes
    train_dfs = []
    test_dfs = []
    
    print(f"Splitting each category (sorted by {sort_by}) without shuffling...")
    for category in expected_categories:
        # Get data for current category and sort
        category_data = data[data['category'] == category].copy()
        category_data = category_data.sort_values(sort_by)
        
        if len(category_data) == 0:
            print(f"Warning: No data for category '{category}'")
            continue
        
        # Calculate split point
        split_idx = int(len(category_data) * train_ratio)
        
        if split_idx == 0 or split_idx == len(category_data):
            print(f"Warning: Category '{category}' has too few samples ({len(category_data)}). "
                  f"Using all for {'train' if split_idx > 0 else 'test'}.")
            if split_idx > 0:
                train_dfs.append(category_data)
            else:
                test_dfs.append(category_data)
            continue
        
        # Split without shuffling: first train_ratio for train, rest for test
        train_cat = category_data.iloc[:split_idx].copy()
        test_cat = category_data.iloc[split_idx:].copy()
        
        print(f"  {category}: {len(train_cat)} train, {len(test_cat)} test "
              f"({len(test_cat)/len(category_data):.1%} test)")
        
        train_dfs.append(train_cat)
        test_dfs.append(test_cat)
    
    # Combine all categories
    print("\nCombining results...")
    
    if train_dfs:
        train_data = pd.concat(train_dfs, ignore_index=False)
    else:
        train_data = pd.DataFrame(columns=data.columns)
        
    if test_dfs:
        test_data = pd.concat(test_dfs, ignore_index=False)
    else:
        test_data = pd.DataFrame(columns=data.columns)
    
    # Remove the temporary category column
    if 'category' in train_data.columns:
        train_data = train_data.drop(columns=['category'])
    if 'category' in test_data.columns:
        test_data = test_data.drop(columns=['category'])
    
    # Save the split datasets
    print("Saving results...")
    train_data.to_parquet(output_train_path, index=False)
    test_data.to_parquet(output_test_path, index=False)
    
    print(f"\nTrain set saved to: {output_train_path}")
    print(f"Test set saved to: {output_test_path}")
    
    return train_data, test_data


INPUT_PATH = "/home/ta-vm/Desktop/ws/ece448labs/spring26/mp02/src/data/train/llm.parquet.0"
OUTPUT_TRAIN_PATH = "/home/ta-vm/Desktop/ws/ece448labs/spring26/mp02/src/data/train/llm.parquet"
OUTPUT_TEST_PATH = "/home/ta-vm/Desktop/ws/ece448labs/spring26/mp02/src/data/train/llm.parquet.1"

# Split without shuffling (preserving original order)
split_ordered_dataset(
    input_path=INPUT_PATH,
    output_train_path=OUTPUT_TRAIN_PATH,
    output_test_path=OUTPUT_TEST_PATH,
    train_ratio=0.18
)