from extremeweatherbench import inputs, metrics, evaluate, utils
# 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'
# Define a forecast object; in this case, a KerchunkForecast
fcnv2_forecast = inputs.KerchunkForecast(
    source=forecast_dir, # source path
    variables=["surface_air_temperature"], # variables to use in the evaluation
    variable_mapping=inputs.CIRA_metadata_variable_mapping, # mapping to use for variables in forecast dataset to EWB variable names
    storage_options={"remote_protocol": "s3", "remote_options": {"anon": True}}, # storage options for access
)
# Load in ERA5; source defaults to the ARCO ERA5 dataset from Google and variable mapping is provided by default as well
era5_heatwave_target = inputs.ERA5(
    variables=["surface_air_temperature"], # variable to use in the evaluation
    storage_options={"remote_options": {"anon": True}}, # storage options for access
    chunks=None, # define chunks for the ERA5 data
)
# EvaluationObjects are used to evaluate a single forecast source against a single target source with a defined event type. Event types are declared with each case. One or more metrics can be evaluated with each EvaluationObject.
heatwave_evaluation_list = [
    inputs.EvaluationObject(
        event_type="heat_wave",
        metric_list=[
            metrics.MaximumMAE,
            metrics.RMSE,
            metrics.OnsetME,
            metrics.DurationME,
            metrics.MaxMinMAE,
        ],
        target=era5_heatwave_target,
        forecast=fcnv2_forecast,
    ),
]
# Load in the EWB default list of event cases
cases = utils.load_events_yaml()
# Create the evaluation class, with cases and evaluation objects declared
ewb_instance = evaluate.ExtremeWeatherBench(
    cases=cases,
    evaluation_objects=heatwave_evaluation_list,
)
# Execute a parallel run and return the evaluation results as a pandas DataFrame
heatwave_outputs = ewb_instance.run(
    n_jobs=16, # use 16 processes
    pre_compute=True, # load case data into memory before metrics are computed. Useful with smaller evaluation datasets with many metrics
)
# Save the results
outputs.to_csv('heatwave_evaluation_results.csv')