from calibrator.metrics import ECE

def test_local_calibrator():
    print("---Test LocalCalibrator---")


    from calibrator import LocalCalibrator
    import torch

    val_logits, val_labels = torch.load("tests/test_logits/resnet50_cifar10_cross_entropy_val_0.1_vanilla.pt", weights_only=False)
    test_logits, test_labels = torch.load("tests/test_logits/resnet50_cifar10_cross_entropy_test_0.9_vanilla.pt", weights_only=False)

    calibrator = LocalCalibrator()
    eps_opt = calibrator.fit(val_logits, val_labels)
    calibrated_probability = calibrator.calibrate(test_logits)

    print(f"eps_opt: {eps_opt:.4f}")
    print(f"Uncalibrated ECE: {ECE()(labels=test_labels, logits=test_logits):.4f}")
    print(f"Calibrated ECE: {ECE()(labels=test_labels, softmaxes=calibrated_probability):.4f}")
    print("!!! Pass LocalCalibrator Test !!!")

if __name__ == "__main__":
    test_local_calibrator()