scripts/collect_sample_data.py

"""
Collect 1-Week Sample Data from JAO
Sept 23-30, 2025 (7 days)

Collects:
- MaxBEX (TARGET VARIABLE)
- Active Constraints (CNECs + PTDFs)
"""

import sys
from pathlib import Path

# Add src to path
sys.path.insert(0, str(Path(__file__).parent.parent / 'src'))

from data_collection.collect_jao import JAOCollector

def main():
    # Initialize collector
    collector = JAOCollector()

    # Define 1-week sample period
    start_date = '2025-09-23'
    end_date = '2025-09-30'

    # Output directory
    output_dir = Path('data/raw/sample')
    output_dir.mkdir(parents=True, exist_ok=True)

    print("\n" + "="*80)
    print("JAO 1-WEEK SAMPLE DATA COLLECTION")
    print("="*80)
    print(f"Period: {start_date} to {end_date} (7 days)")
    print(f"Output: {output_dir}")
    print("="*80 + "\n")

    # Collect MaxBEX (TARGET)
    maxbex_path = output_dir / 'maxbex_sample_sept2025.parquet'
    print("\n[1/2] Collecting MaxBEX (TARGET VARIABLE)...")
    print("Estimated time: ~35 seconds (7 days × 5 sec rate limit)\n")

    maxbex_df = collector.collect_maxbex_sample(
        start_date=start_date,
        end_date=end_date,
        output_path=maxbex_path
    )

    # Collect CNECs + PTDFs
    cnec_path = output_dir / 'cnecs_sample_sept2025.parquet'
    print("\n[2/2] Collecting Active Constraints (CNECs + PTDFs)...")
    print("Estimated time: ~35 seconds (7 days × 5 sec rate limit)\n")

    cnec_df = collector.collect_cnec_ptdf_sample(
        start_date=start_date,
        end_date=end_date,
        output_path=cnec_path
    )

    # Summary
    print("\n" + "="*80)
    print("SAMPLE DATA COLLECTION COMPLETE")
    print("="*80)

    if maxbex_df is not None:
        print(f"[OK] MaxBEX: {maxbex_path}")
        print(f"     Shape: {maxbex_df.shape}")
    else:
        print("[ERROR] MaxBEX collection failed")

    if cnec_df is not None:
        print(f"[OK] CNECs/PTDFs: {cnec_path}")
        print(f"     Shape: {cnec_df.shape}")
    else:
        print("[ERROR] CNEC/PTDF collection failed")

    print("\nNext step: Run Marimo notebook for data exploration")
    print("Command: marimo edit notebooks/01_data_exploration.py")
    print("="*80 + "\n")

if __name__ == '__main__':
    main()