# Finetuning a selected model 
# on a selected dataset 
# using selected parameters

from torch.utils.data import DataLoader
from sklearn.metrics import classification_report
import SetUp
from Data import *
#from datasets.Casia import *
#from IdentitySubset import IdentitySubset
from datasets.UniversalIdentitySubset import UniversalIdentitySubset as IdentitySubset
# models
from architectures.Model import Model, Architecture

from unlearning import LinearFiltration, WeightFiltration, CertifiedRemoval

# numbre of classes 
CLASS_SIZE = 20
# batch 
BATCH_SIZE = 32

# size of images per class trainset + testset
# 30 works best, more than that and we dont have enough data
SAMPLE_SIZE = 30

# this is then (full_sample - test_sample)
TRAINING_SMPLE = 28

# learning rate
LR_RATE = 0.0001
EPOCHS = 20

# depends on model architecture
# ResNet, DenseNet = 224
# Inception = 299
RESOLUTION = 224

# model architecture options are 
#  - RESNET18
#  - RESNET50
#  - DENSENET121
#  - INCEPTION
#  - GOOGLENET
#  - EFFICIENTNET
#  - SHUFFLENET
arch = Architecture.RESNET50

# DATA PREPARATION
# load data set and prepare
dataset = get_set()
# select identities for experiment
#selected_identities = select_ids(
#    dataset = dataset,
#    sample_size = SAMPLE_SIZE,
#    class_size = CLASS_SIZE
#    )

# this selects the top 50 based on sample size
# that way repeated calls return the same classes
selected_identities = select_top_ids(
    dataset=dataset,
    class_size= CLASS_SIZE,
)

print(f'> Selected {CLASS_SIZE} random identity classes from CelebA dataset.')
print(f'> A class has {TRAINING_SMPLE} train and {SAMPLE_SIZE-TRAINING_SMPLE} test samples')

# split class images to train/test indices
train_indices, test_indices = get_indices(
    dataset = dataset, 
    identities = selected_identities, 
    split_at = TRAINING_SMPLE
)

# helps map class id to index
id_map = {old_id: new_id for new_id, old_id in enumerate(selected_identities)}

# we remap identities because crossEntropyLoss requires in indices 0 -> (n-1)
# where n = class size.
tr_transform = train_transform(res = RESOLUTION)
train_data = IdentitySubset(
    dataset=dataset, 
    indices=train_indices, 
    id_mapping=id_map,
    transform=tr_transform)

train_loader = DataLoader(
    train_data, 
    batch_size = BATCH_SIZE, 
    shuffle = True)

print(f"> Total training images: {len(train_data)}")

print(f'> Constants : Classes = {CLASS_SIZE}, Batch = {BATCH_SIZE}, epochs = {EPOCHS}')

# MODEL PREPARATION
# cuda if exists (it does here)
device = SetUp.get_device()
# Create model using Factory
model = Model.create(
    arch = arch, 
    device = device, 
    size = CLASS_SIZE)

# we may need to load existing model or finetune
model.train(
        epochs = EPOCHS, 
        loader = train_loader, 
        rate = LR_RATE)

    # save.
model.save(filename=arch.name.lower())


# done tuning
print('Model saved!')

# EVALUATE

te_transform = test_transform(RESOLUTION)
# Testing
test_data = IdentitySubset(
    dataset = dataset, 
    indices=test_indices, 
    id_mapping=id_map,
    transform=te_transform)

test_loader = DataLoader(
    test_data, 
    batch_size=BATCH_SIZE, 
    shuffle=False)

print(f"Total test images for these {CLASS_SIZE} classes: {len(test_data)}")

# Evaluate
model.evaluate(
    loader = test_loader)

# test again
reloaded = Model.create(
    arch=arch,
    device = device,
    size = CLASS_SIZE
    )
reloaded.load(arch = arch)
print("Evaluating loaded")
reloaded.evaluate(
    loader = test_loader
)


strategies_to_test = [
    LinearFiltration(target_class_idx=12),
    WeightFiltration(target_class_idx=12),
    CertifiedRemoval(target_class_idx=12)
]

# Run the comparative benchmark seamlessly
execution_profiles = {}
for strategy in strategies_to_test:
    # Each iteration clones weights back to fine-tuned state before running
    runtime = my_model.unlearn(strategy, forget_loader, retain_loader)
    execution_profiles[strategy.__class__.__name__] = runtime