import pickle
import os

SOURCE = "energizer-grid-data/full_dataset.pkl"
DEST = "data/sample_dataset.pkl"

def main():
    if not os.path.exists(SOURCE):
        print(f"Error: Source {SOURCE} not found!")
        return

    print(f"Reading {SOURCE}...")
    with open(SOURCE, "rb") as f:
        full_data = pickle.load(f)

    print(f"Dataset contains {len(full_data)} days.")
    
    # Extract first 7 days
    sample_data = full_data[:7]
    print(f"Extracted {len(sample_data)} days.")

    # Ensure dest dir exists
    os.makedirs(os.path.dirname(DEST), exist_ok=True)

    with open(DEST, "wb") as f:
        pickle.dump(sample_data, f)
    
    print(f"Success! Sample saved to {DEST}")

if __name__ == "__main__":
    main()