from extremeweatherbench import config, events, evaluate
import pickle 
# Select model
model = 'FOUR_v200_GFS'
# Set up path to directory of file - zarr or kerchunk/virtualizarr json/parquet
forecast_dir = f'gs://extremeweatherbench/{model}.parq'
# Choose the event types you want to include
event_list = [events.HeatWave,
              events.Freeze]
# Use ForecastSchemaConfig to map forecast variable names to CF convention-based names used in EWB
# the sample forecast kerchunk references to the CIRA MLWP archive are the default configuration
default_forecast_config = config.ForecastSchemaConfig()
# Set up configuration object that includes events and the forecast directory
heatwave_and_freeze_configuration = config.Config(
    event_types=event_list,
    forecast_dir=forecast_dir,
    # This line is not necessary, forecast_schema_config defaults to the default_forecast_config.
    # Here as an example if values need to be changed for your use case 
    forecast_schema_config=default_forecast_config 
    )
# Run the evaluate script which outputs a dataframe of case results with associated metrics and variables
cases = evaluate.evaluate(eval_config=heatwave_and_freeze_configuration)
# Save the results to a pickle file
with open(f'ewb_cases_{model}.pkl', 'wb') as f:
    pickle.dump(cases, f)
# Or, save to csv:
cases.to_csv(f'ewb_cases_{model}.csv')