import pandas as pd
import numpy as np
import os

BASE_DIR = "/opt/airflow/process"
INPUT_FILE = os.path.join(
    BASE_DIR, "Extended_Employee_Performance_and_Productivity_Data.csv")
BATCH_DIR = os.path.join(BASE_DIR, "batches")

df = pd.read_csv(INPUT_FILE)
os.makedirs(BATCH_DIR, exist_ok=True)

batches = np.array_split(df, 10)

for i, batch in enumerate(batches, 1):
    file_name = os.path.join(BATCH_DIR, f"batch_{i}.csv")
    batch.to_csv(file_name, index=False)
    print(f"Saved: {file_name}")