import os
import pandas as pd
import glob

# Абсолютный путь внутри контейнера
batch_folder = "/opt/airflow/batches_for_5years"
os.makedirs(batch_folder, exist_ok=True)

source_file = "/opt/airflow/data/us_rates.csv"
df_source = pd.read_csv(source_file, parse_dates=['Date'])

batch_years = 5
start_year = df_source['Date'].dt.year.min()
end_year = df_source['Date'].dt.year.max()

existing_batches = glob.glob(os.path.join(batch_folder, "us_rates_*.csv"))
existing_numbers = [int(os.path.splitext(os.path.basename(f))[0].split("_")[-1])
                    for f in existing_batches if os.path.splitext(os.path.basename(f))[0].split("_")[-1].isdigit()]

next_number = 1 if not existing_numbers else max(existing_numbers) + 1

for batch_start in range(start_year, end_year + 1, batch_years):
    batch_end = min(batch_start + batch_years - 1, end_year)
    df_batch = df_source[(df_source['Date'].dt.year >= batch_start) & (
        df_source['Date'].dt.year <= batch_end)]

    if not df_batch.empty:
        batch_file = os.path.join(batch_folder, f"us_rates_{next_number}.csv")
        df_batch.to_csv(batch_file, index=False)
        print(f"Batch {next_number} saved: {batch_file}")
        next_number += 1