import threshold_xarray as th
import xarray as xr
import numpy as np
import torch
import os

import sklearn.calibration as cal


# The purpose of this file is to calibrate the data on the test set, and evaluate the calibration on the validation set.
# We are using scikits calibration library to do this.

if __name__ == '__main__':
    print('Loading Config..B')
    config = th.load_config()
    ENSEMBLE_PATH = f"{config['paths']['model_output']}{config['ensemble']['name']}"
    V4_PATH = ENSEMBLE_PATH + '/v4'

    if not os.path.exists(V4_PATH):
        os.makedirs(V4_PATH)
    print('Config Loaded')

    # Load the predictions
    print('Loading Predictions...')
    val_preds = xr.open_dataset(f'{ENSEMBLE_PATH}/val_predictions.nc')
    test_preds = xr.open_dataset(f'{ENSEMBLE_PATH}/test_predictions.nc')
    print('Predictions Loaded')

    # Now the goal is to calibrate the test set, and evaluate the calibration on the validation set.
    # We do this by binning the data into 15 bins, and then calculating the mean of the predictions in each bin.
    # We then use this to calibrate the data.

    # First, get the statistics of both sets
    print('Calculating Statistics...')
    val_stats = th.compute_ensemble_statistics(val_preds)
    test_stats = th.compute_ensemble_statistics(test_preds)

    # Calibrate the test set
    print('Calibrating Test Set...')