"""
Quick demonstration of XyzDataset streaming functionality.
"""
import tempfile
from chem_mat_data import XyzDataset

print("=" * 60)
print("XyzDataset Demonstration")
print("=" * 60)

# Create XyzDataset with streaming access to XYZ files
with tempfile.TemporaryDirectory() as temp_path:
    dataset = XyzDataset(
        dataset='_test2',
        parser_cls='default',
        folder_path=temp_path,
        use_cache=True,
        target_columns=None  # No properties for this demo
    )

    print(f"\nDataset: {dataset.dataset}")
    print(f"Total molecules: {len(dataset)}")
    print(f"Parser class: {dataset.parser_cls}")

    # Iterate through first 3 molecules
    print("\n" + "-" * 60)
    print("Streaming through first 3 molecules:")
    print("-" * 60)

    for i, (xyz_data, properties) in enumerate(dataset):
        if i >= 3:
            break

        print(f"\nMolecule {i+1}:")
        print(f"  Number of atoms: {xyz_data['num_atoms']}")
        print(f"  Atomic numbers: {xyz_data['atomic_numbers']}")
        print(f"  Atom symbols: {xyz_data['symbols']}")
        print(f"  Positions shape: {xyz_data['positions'].shape}")
        print(f"  Position sample (first atom): {xyz_data['positions'][0]}")
        print(f"  Properties: {properties}")

print("\n" + "=" * 60)
print("✓ XyzDataset streaming works perfectly!")
print("=" * 60)
