# Import the robust deep learning library as rdl
import robust_deep_learning as rdl

#####################################################################################################
#####################################################################################################
# Training or Fine-tuning
#####################################################################################################
#####################################################################################################

#########################################################
# Option 1: Creating a custom model and defining the loss
#########################################################

# Create from a model definition file. 
# For example, import a class "Model" from a model definition file.
model = Model()
# Load a pretrained model if fine-tuning rather than training from scratch (random weights).
# model.load_state_dict(torch.load(pre_trained_file_name, map_location="cuda:" + str(args.gpu)))

# Chance the output last layer of the model with the desired loss first part.
# If the name of the output last layer of the model is unknown, print it with "print(model)".
# For example, if the output last layer is called "classifier":
# "model.classifier = nn.Linear(num_features, num_classes)"
# Then, replace this output layer (usually a linear layer) by adding the following line: 
model.classifier = rdl.DisMaxLossFirstPart(num_features, num_classes)

# Replace the Cross-Entropy Loss.
# The first argument is the name of the output last layer of the model used above.
# In case of training using a not too constrained model on image data, pass an add-on.
# Currently, only DisMax has one add-on called "FPR".
# Otherwise, do not pass add-on or simple pass "add_on=None". 
criterion = rdl.DisMaxLossSecondPart(model.classifier, add_on="FPR", gpu=None)

#######################################################
# Option 2: Creating a timm model and defining the loss
#######################################################

# For using a model from timm lib, use "create_model" functionality.
# It is possible to start from a pretrained model to fine-tune using the desired loss.
model = timm.create_model('resnet18', pretrained=False)

# Chance the output last layer of the model with the desired loss first part.
# If the name of the output last layer of the model is unknown, print it with "print(model)".
# For example, if the output last layer is called "fc":
# "model.fc = nn.Linear(num_features, num_classes)"
# Then, replace this output layer (usually a linear layer) by adding the following line: 
model.fc = rdl.DisMaxLossFirstPart(model.get_classifier().in_features, num_classes)

# Replace the Cross-Entropy Loss.
# The first argument is the name of the output last layer of the model used above.
# In case of training using a not too constrained model on image data, pass an add-on.
# Currently, only DisMax has one add-on called "FPR".
# Otherwise, do not pass add-on or simple pass "add_on=None". 
criterion = rdl.DisMaxLossSecondPart(model.fc, add_on="FPR", gpu=None)

############################
# Checking the training loop
############################

for epoch in epochs:

    # Training loop
    for inputs, targets in in_data_train_loader:

        # In the training loop, add the line below for preprocessing before forwarding.
        # This is only required if using add_on other than None. Otherwise, this line is not needed.
        inputs, targets = criterion.preprocess(inputs, targets) 

        # The three below lines are usually found in the training loop!
        # These lines must not be changed!
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()

#####################################################################################################
#####################################################################################################
# Uncertainty Estimation
#####################################################################################################
#####################################################################################################

#############
# Calibrating
#############

# Get outputs, labels, accuracy, expected calibration error (ECE), and negative log-likelihood (NLL).
results = rdl.get_outputs_labels_and_metrics(model, in_data_val_loader, gpu=None)

# Calculate probabilities and verify the values returned above. 
probabilities = torch.nn.Softmax(dim=1)(results["outputs"])
print(probabilities)
print(results["acc"], results["ece"], results["nll"])

# Calibrate the temperature passing the output last layer, the model, and the validation set.
# Choose the metric to optimize. For example, "ECE" (the only option for now).
rdl.calibrate_temperature(model.classifier, model, in_data_val_loader, optimize="ECE", gpu=None)
#rdl.calibrate_temperature(model.fc, model, in_data_val_loader, optimize="ECE")

# Verify the novel value of the temperature of the output last layer.
print(model.classifier.temperature)

######################
# Verifing calibration
######################

# Get the results again using the calibrated model.
results = rdl.get_outputs_labels_and_metrics(model, in_data_val_loader, gpu=None)

# Verifiy the probabilities, ECE, and NLL are improved regarding the now calibrated model.
probabilities = torch.nn.Softmax(dim=1)(results["outputs"])
print(probabilities)
print(results["acc"], results["ece"], results["nll"])

####################################################################################################
####################################################################################################
# Out-of-Distribution Detection
####################################################################################################
####################################################################################################

########################
# Estimating performance
########################

# Define a score type. Typically, the best/recommended option for the loss you are using.
score_type = "MMLES"

# Evaluate the out-of-distribution detection performance passing loaders.
ood_metrics = rdl.get_ood_metrics(
    model, in_data_val_loader, out_data_loader, score_type, fpr=0.05, gpu=None)

# Optionally, first get in-data scores:
results = rdl.get_outputs_labels_and_metrics(model, in_data_val_loader, gpu=None)
in_data_scores = rdl.get_scores(results["outputs"], score_type)

# Second, get out-data scores:
results = rdl.get_outputs_labels_and_metrics(model, out_data_loader, gpu=None)
out_data_scores = rdl.get_scores(results["outputs"], score_type)

# Then, finally, evaluate the out-of-distribution detection performance passing scores.
ood_metrics = rdl.get_ood_metrics_from_scores(in_data_scores, out_data_scores, fpr=0.05)

######################################
# Detecting (still testing this part))
######################################

# Before detecting, it is necessary to compute the thresholds.
thresholds = rdl.get_thresholds(model, in_data_val_loader, score_type, gpu=None)

# Optionaly, it is possible to compute the threshold using in-data scores.
#thresholds = rdl.get_thresholds(results["outputs"], score_types="MMLES")
thresholds = rdl.get_thresholds_from_scores(in_data_scores) # guarder 

# After calculating thresholds, detection may be performed.
# Some test input may be obtained using: input = next(iter(in_data_val_loader))
ood_detections = rdl.get_ood_detections(model, inputs, thresholds, fpr="0.05", gpu=None)