from femb.backbones import build_backbone
from femb.headers import ArcFaceHeader
from femb import FaceEmbeddingModel

# build the backbone embedding network
backbone = build_backbone(backbone="iresnet18", embed_dim=embed_dim)

# create one of the face recognition headers
# header = ArcFaceHeader(in_features=embed_dim, out_features=train_n_classes)
header = MagFaceHeader(in_features=embed_dim, out_features=train_n_classes)

# create the ce loss
loss = torch.nn.CrossEntropyLoss()

# create the face recognition model wrapper
face_model = FaceEmbeddingModel(backbone=backbone, header=header, loss=loss)

from femb.evaluation import VerificationEvaluator

# create the verification evaluator
evaluator = VerificationEvaluator(similarity='cos')

# specify the optimizer (and a scheduler)
optimizer = torch.optim.SGD(params=face_model.params, lr=1e-2, momentum=0.9, weight_decay=5e-4)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer=optimizer, milestones=[8000, 10000, 160000], gamma=0.1)

# fit the face embedding model to the dataset
face_model.fit(
    train_dataset=train_dataset,        # specify the training set
    batch_size=32,                      # batch size for training and evaluation
    device='cuda',                      # torch device, i.e. 'cpu' or 'cuda'
    optimizer=optimizer,                # torch optimizer
    lr_epoch_scheduler=None,            # scheduler based on epochs
    lr_global_step_scheduler=scheduler, # scheduler based on global steps
    evaluator=evaluator,                # evaluator module
    val_dataset=val_dataset,            # specify the validation set
    evaluation_steps=10,                # number of steps between evaluations
    max_training_steps=20000,           # maximum number of (global) training steps (if zero then max_epochs count is used for stopping)
    max_epochs=0,                       # maximum number of epochs (if zero then max_training_steps is used for stopping)
    tensorboard=True                    # specify whether or not tensorboard shall be used for embedding projections and metric monitoring
    )