import argparse
import pandas as pd
from sklearn.model_selection import train_test_split
import os



def parse(csv_path):
    print(f"Location of the file: {csv_path}")
    
    # Step 1: Load the dataset
    # file_path = "dataset.csv"  # Path to the original dataset
    data = pd.read_csv(csv_path)

    # Drop dupes
    data = data.drop_duplicates()
    
    # Step 2: Define the feature columns (X) and target column (y)
    X = data[["name", "attendance percentage", "average sleep time", "average screen time"]]  # Feature columns
    y = data["grade"]  # Target column
    
    # Step 3: Split the dataset into training and testing sets
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    # Step 4: Combine X and y back into dataframes for train and test
    train_data = pd.concat([X_train, y_train], axis=1)  # Combine features and target for training data
    test_data = pd.concat([X_test, y_test], axis=1)    # Combine features and target for testing data
    
    # Step 5: Create the 'data' folder if it doesn't exist
    output_folder = "data"
    os.makedirs(output_folder, exist_ok=True)
    
    # Step 6: Save the train and test sets as CSV files
    train_file_path = os.path.join(output_folder, "train.csv")
    test_file_path = os.path.join(output_folder, "test.csv")
    
    train_data.to_csv(train_file_path, index=False)
    test_data.to_csv(test_file_path, index=False)
    
    print(f"Train and test datasets saved in '{output_folder}' folder.")
    




if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--csv-path", type=str)
    
    
    args = parser.parse_args()
    parse(args.csv_path)