Blueprint Catalog API

%load_ext autoreload
%autoreload 2

from cson_forge import catalog

# Find all blueprint files (defaults to all stages)
blueprint_files = catalog.blueprint.find_blueprint_files()
print(f"Found {len(blueprint_files)} blueprint files:")
for bp_file in blueprint_files[:10]:  # Show first 10
    print(f"  - {bp_file.name}")

# You can also filter by stage
postconfig_files = catalog.blueprint.find_blueprint_files(stage="postconfig")
print(f"\nFound {len(postconfig_files)} postconfig blueprint files")

# Load a single blueprint
if blueprint_files:
    bp_data = catalog.blueprint.load_blueprint(blueprint_files[0])
    blueprint_name = bp_data.get('name', '')
    model_name, grid_name = catalog.blueprint._extract_model_and_grid_name(blueprint_name)
    partitioning = bp_data.get('partitioning', {})
    
    print(f"Blueprint name: {blueprint_name}")
    print(f"Model name: {model_name}")
    print(f"Grid name: {grid_name}")
    print(f"Description: {bp_data.get('description')}")
    print(f"Start time: {bp_data.get('valid_start_date')}")
    print(f"End time: {bp_data.get('valid_end_date')}")
    if isinstance(partitioning, dict):
        print(f"Processors: {partitioning.get('n_procs_x')} x {partitioning.get('n_procs_y')}")

# Load grid kwargs from a blueprint
# Grid YAML files are typically in the same directory as the blueprint
if blueprint_files:
    bp_file = blueprint_files[0]
    bp_data = catalog.blueprint.load_blueprint(bp_file)

    # Look for _grid.yml in the same directory as the blueprint
    grid_yaml_path = bp_file.parent / "_grid.yml"
    
    if grid_yaml_path.exists():
        try:
            grid_kwargs = catalog.blueprint.load_grid_kwargs(grid_yaml_path)
            print("Grid parameters:")                
            for key, value in grid_kwargs.items():
                print(f"  {key}: {value}")
        except Exception as e:
            print(f"Could not load grid kwargs: {e}")
    else:
        print(f"Grid YAML file not found at {grid_yaml_path}")
        print("Grid parameters may be available in the DataFrame after calling load()")

# Load all blueprints into a DataFrame
# Defaults to 'postconfig' stage which has the most complete data
df = catalog.blueprint.load(stage="postconfig")

print(f"Loaded {len(df)} blueprints")
print(f"\nDataFrame columns: {list(df.columns)}")
print(f"\nDataFrame shape: {df.shape}")

# Display the DataFrame (excluding dict columns for readability)
display_cols = [col for col in df.columns if col not in ['grid_kwargs']]
df[display_cols]

# Display basic information about the DataFrame
if not df.empty:
    print("DataFrame info:")
    print(df.info())
    
    print("\nFirst few rows:")
    # Display non-dict columns for readability
    display_cols = [col for col in df.columns if col not in ['grid_kwargs']]
    display(df[display_cols].head())

# Display grid kwargs for the first blueprint
if not df.empty:
    first_row = df.iloc[0]
    print(f"Grid kwargs for {first_row['grid_name']}:")
    grid_kwargs = first_row['grid_kwargs']
    if isinstance(grid_kwargs, dict):
        for key, value in grid_kwargs.items():
            print(f"  {key}: {value}")

# Query by model name
if not df.empty:
    model_name = df['model_name'].iloc[0] if 'model_name' in df.columns else None
    if model_name:
        model_blueprints = df[df['model_name'] == model_name]
        print(f"Found {len(model_blueprints)} blueprints for model '{model_name}':")
        print(model_blueprints[['grid_name', 'start_time', 'end_time']].to_string())

# Query by grid name
if not df.empty and 'grid_name' in df.columns:
    grid_name = df['grid_name'].iloc[0]
    grid_blueprints = df[df['grid_name'] == grid_name]
    print(f"\nFound {len(grid_blueprints)} blueprints for grid '{grid_name}':")
    print(grid_blueprints[['model_name', 'start_time', 'end_time']].to_string())

# Example: Using blueprint data to create a CstarSpecBuilder
from cson_forge import CstarSpecBuilder
from datetime import datetime

# Example: Create a builder from blueprint data
if not df.empty:
    row = df.iloc[0]
    
    # Note: To recreate a builder from a blueprint, you would need to:
    # 1. Load the blueprint file using CstarSpecBuilder.from_blueprint() (if such method exists)
    # 2. Or manually extract data from the DataFrame and create a new builder
    # 3. The blueprint_path column contains the path to the blueprint file
    
    print(f"Blueprint path: {row['blueprint_path']}")
    print(f"Model: {row['model_name']}, Grid: {row['grid_name']}")
    print(f"Time range: {row['start_time']} to {row['end_time']}")
    
    # TODO: Add method to CstarSpecBuilder to load from existing blueprint file