"""

train_classifier.py

Train a EfficientNet or ResNet classifier.

Currently the implementation of multi-label multi-class classification is
non-functional.

During training, start tensorboard from within the classification/ directory:
  tensorboard --logdir run --bind_all --samples_per_plugin scalars=0,images=0

"""

#%% Imports and constants

from __future__ import annotations

import json
import os
import argparse

from collections import defaultdict
from collections.abc import Callable, Mapping, MutableMapping, Sequence
from datetime import datetime
from typing import Any

import numpy as np
import PIL.Image
import sklearn.metrics
import tqdm

import torch
from torch.utils import tensorboard
import torchvision as tv
from torchvision.datasets.folder import default_loader

from megadetector.classification import efficientnet, evaluate_model
from megadetector.classification.train_utils import (
    HeapItem, recall_from_confusion_matrix, add_to_heap, fig_to_img,
    imgs_with_confidences, load_dataset_csv, prefix_all_keys)
from megadetector.visualization import plot_utils


#%% Example usage

"""
    python train_classifier.py run_idfg /ssd/crops_sq \
        -m "efficientnet-b0" --pretrained --finetune --label-weighted \
        --epochs 50 --batch-size 512 --lr 1e-4 \
        --num-workers 12 --seed 123 \
        --logdir run_idfg
"""

# mean/std values from https://pytorch.org/docs/stable/torchvision/models.html
MEANS = np.asarray([0.485, 0.456, 0.406])
STDS = np.asarray([0.229, 0.224, 0.225])

VALID_MODELS = sorted(
    set(efficientnet.VALID_MODELS) |
    {'resnet101', 'resnet152', 'resnet18', 'resnet34', 'resnet50'})


class AverageMeter:
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self) -> None:
        self.val = 0.0
        self.avg = 0.0
        self.sum = 0.0
        self.count = 0

    def update(self, val: float, n: int = 1) -> None:
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


class SimpleDataset(torch.utils.data.Dataset):
    """A simple dataset that simply returns images and labels."""

    def __init__(self,
                 img_files: Sequence[str],
                 labels: Sequence[Any],
                 sample_weights: Sequence[float] | None = None,
                 img_base_dir: str = '',
                 transform: Callable[[PIL.Image.Image], Any] | None = None,
                 target_transform: Callable[[Any], Any] | None = None):
        """Creates a SimpleDataset."""
        self.img_files = img_files
        self.labels = labels
        self.sample_weights = sample_weights
        self.img_base_dir = img_base_dir
        self.transform = transform
        self.target_transform = target_transform

        self.len = len(img_files)
        assert len(labels) == self.len
        if sample_weights is not None:
            assert len(sample_weights) == self.len

    def __getitem__(self, index: int) -> tuple[Any, ...]:
        """
        Args:
            index: int

        Returns: tuple, (sample, target) or (sample, target, sample_weight)
        """
        img_file = self.img_files[index]
        img = default_loader(os.path.join(self.img_base_dir, img_file))
        if self.transform is not None:
            img = self.transform(img)
        target = self.labels[index]
        if self.target_transform is not None:
            target = self.target_transform(target)
        if self.sample_weights is not None:
            return img, target, img_file, self.sample_weights[index]
        return img, target, img_file

    def __len__(self) -> int:
        return self.len


def create_dataloaders(
        dataset_csv_path: str,
        label_index_json_path: str,
        splits_json_path: str,
        cropped_images_dir: str,
        img_size: int,
        multilabel: bool,
        label_weighted: bool,
        weight_by_detection_conf: bool | str,
        batch_size: int,
        num_workers: int,
        augment_train: bool
        ) -> tuple[dict[str, torch.utils.data.DataLoader], list[str]]:
    """
    Args:
        dataset_csv_path: str, path to CSV file with columns
            ['dataset', 'location', 'label'], where label is a comma-delimited
            list of labels
        splits_json_path: str, path to JSON file
        augment_train: bool, whether to shuffle/augment the training set

    Returns:
        datasets: dict, maps split to DataLoader
        label_names: list of str, label names in order of label id
    """
    df, label_names, split_to_locs = load_dataset_csv(
        dataset_csv_path, label_index_json_path, splits_json_path,
        multilabel=multilabel, label_weighted=label_weighted,
        weight_by_detection_conf=weight_by_detection_conf)

    # define the transforms
    normalize = tv.transforms.Normalize(mean=MEANS, std=STDS, inplace=True)
    train_transform = tv.transforms.Compose([
        tv.transforms.RandomResizedCrop(img_size),
        tv.transforms.RandomRotation(degrees=(-90, 90)),
        tv.transforms.RandomHorizontalFlip(p=0.5),
        tv.transforms.RandomVerticalFlip(p=0.1),
        tv.transforms.RandomGrayscale(p=0.1),
        tv.transforms.ColorJitter(brightness=.25, contrast=.25, saturation=.25),
        tv.transforms.ToTensor(),
        normalize
    ])
    test_transform = tv.transforms.Compose([
        # resizes smaller edge to img_size
        tv.transforms.Resize(img_size, interpolation=PIL.Image.BICUBIC),
        tv.transforms.CenterCrop(img_size),
        tv.transforms.ToTensor(),
        normalize
    ])

    dataloaders = {}
    for split, locs in split_to_locs.items():
        is_train = (split == 'train') and augment_train
        split_df = df[df['dataset_location'].isin(locs)]

        sampler: torch.utils.data.Sampler | None = None
        weights = None
        if label_weighted or weight_by_detection_conf:
            # weights sums to:
            # - if weight_by_detection_conf: (# images in split - conf delta)
            # - otherwise: # images in split
            weights = split_df['weights'].to_numpy()
            if not weight_by_detection_conf:
                assert np.isclose(weights.sum(), len(split_df))
            if is_train:
                sampler = torch.utils.data.WeightedRandomSampler(
                    weights, num_samples=len(split_df), replacement=True)
        elif is_train:
            # for normal (non-weighted) shuffling
            sampler = torch.utils.data.SubsetRandomSampler(range(len(split_df)))

        dataset = SimpleDataset(
            img_files=split_df['path'].tolist(),
            labels=split_df['label_index'].tolist(),
            sample_weights=weights,
            img_base_dir=cropped_images_dir,
            transform=train_transform if is_train else test_transform)
        assert len(dataset) > 0
        dataloaders[split] = torch.utils.data.DataLoader(
            dataset, batch_size=batch_size, sampler=sampler,
            num_workers=num_workers, pin_memory=True)

    return dataloaders, label_names


def set_finetune(model: torch.nn.Module, model_name: str, finetune: bool
                 ) -> None:
    """Set the 'requires_grad' on each model parameter according to whether or
    not we are fine-tuning the model.
    """
    if finetune:
        if 'efficientnet' in model_name:
            final_layer = model._fc  # pylint: disable=protected-access
        else:  # torchvision resnet
            final_layer = model.fc
        assert isinstance(final_layer, torch.nn.Module)

        # set all parameters to not require gradients except final FC layer
        model.requires_grad_(False)
        for param in final_layer.parameters():
            param.requires_grad = True
    else:
        model.requires_grad_(True)


def build_model(model_name: str, num_classes: int, pretrained: bool | str,
                finetune: bool) -> torch.nn.Module:
    """Creates a model with an EfficientNet or ResNet base. The model outputs
    unnormalized logits.

    Args:
        model_name: str, name of EfficientNet or Resnet model
        num_classes: int, number of classes for output layer
        pretrained: bool or str, (bool) whether to initialize to ImageNet
            weights, (str) path to checkpoint
        finetune: bool, whether to freeze all layers except the final FC layer

    Returns: torch.nn.Module, model loaded on CPU
    """
    assert model_name in VALID_MODELS

    if 'efficientnet' in model_name:
        if pretrained is True:
            model = efficientnet.EfficientNet.from_pretrained(
                model_name, num_classes=num_classes)
        else:
            model = efficientnet.EfficientNet.from_name(
                model_name, num_classes=num_classes)
    else:
        model_class = getattr(tv.models, model_name)
        model = model_class(pretrained=(pretrained is True))

        # replace final fully-connected layer (which has 1000 ImageNet classes)
        model.fc = torch.nn.Linear(model.fc.in_features, num_classes)

    if isinstance(pretrained, str):
        print(f'Loading saved weights from {pretrained}')
        ckpt = torch.load(pretrained, map_location='cpu')
        model.load_state_dict(ckpt['model'])

    assert all(p.requires_grad for p in model.parameters())
    set_finetune(model=model, model_name=model_name, finetune=finetune)
    return model


def prep_device(model: torch.nn.Module, device_id: int | None = None
                ) -> tuple[torch.nn.Module, torch.device]:
    """Place model on appropriate device.

    Args:
        model: torch.nn.Module, not already wrapped with DataParallel
        device_id: optional int, GPU device to use
            if None, then uses DataParallel when possible
            if specified, then only uses specified device

    Returns:
        model: torch.nn.Module, model placed on <device>, wrapped with
            DataParallel if more than 1 GPU is found
        device: torch.device, 'cuda:{device_id}' if GPU is found, otherwise 'cpu'
    """
    # detect GPU, use all if available
    if torch.cuda.is_available():
        print('CUDA available')
        torch.backends.cudnn.benchmark = True
        if device_id is not None:
            print(f'Starting CUDA device {device_id}')
            device = torch.device(f'cuda:{device_id}')
        else:
            device = torch.device('cuda:0')
            device_ids = list(range(torch.cuda.device_count()))
            if len(device_ids) > 1:
                print(f'Found multiple devices, enabling data parallelism ({device_ids})')
                model = torch.nn.DataParallel(model, device_ids=device_ids)
    else:
        print('CUDA not available, running on the CPU')
        device = torch.device('cpu')
    model.to(device)  # in-place
    return model, device


def main(dataset_dir: str,
         cropped_images_dir: str,
         multilabel: bool,
         model_name: str,
         pretrained: bool | str,
         finetune: int,
         label_weighted: bool,
         weight_by_detection_conf: bool | str,
         epochs: int,
         batch_size: int,
         lr: float,
         weight_decay: float,
         num_workers: int,
         logdir: str,
         log_extreme_examples: int,
         seed: int | None = None) -> None:
    """Main function."""
    # input validation
    assert os.path.exists(dataset_dir)
    assert os.path.exists(cropped_images_dir)
    if isinstance(weight_by_detection_conf, str):
        assert os.path.exists(weight_by_detection_conf)
    if isinstance(pretrained, str):
        assert os.path.exists(pretrained)

    # set seed
    seed = np.random.randint(10_000) if seed is None else seed
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    # create logdir and save params
    params = dict(locals())  # make a copy
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')  # '20200722_110816'
    logdir = os.path.join(logdir, timestamp)
    os.makedirs(logdir, exist_ok=True)
    print('Created logdir:', logdir)
    params_json_path = os.path.join(logdir, 'params.json')
    with open(params_json_path, 'w') as f:
        json.dump(params, f, indent=1)

    if 'efficientnet' in model_name:
        img_size = efficientnet.EfficientNet.get_image_size(model_name)
    else:
        img_size = 224

    # create dataloaders and log the index_to_label mapping
    print('Creating dataloaders')
    loaders, label_names = create_dataloaders(
        dataset_csv_path=os.path.join(dataset_dir, 'classification_ds.csv'),
        label_index_json_path=os.path.join(dataset_dir, 'label_index.json'),
        splits_json_path=os.path.join(dataset_dir, 'splits.json'),
        cropped_images_dir=cropped_images_dir,
        img_size=img_size,
        multilabel=multilabel,
        label_weighted=label_weighted,
        weight_by_detection_conf=weight_by_detection_conf,
        batch_size=batch_size,
        num_workers=num_workers,
        augment_train=True)

    writer = tensorboard.SummaryWriter(logdir)

    # create model
    model = build_model(model_name, num_classes=len(label_names),
                        pretrained=pretrained, finetune=finetune > 0)
    model, device = prep_device(model)

    # define loss function and optimizer
    loss_fn: torch.nn.Module
    if multilabel:
        loss_fn = torch.nn.BCEWithLogitsLoss(reduction='none').to(device)
    else:
        loss_fn = torch.nn.CrossEntropyLoss(reduction='none').to(device)

    # using EfficientNet training defaults
    # - batch norm momentum: 0.99
    # - optimizer: RMSProp, decay 0.9 and momentum 0.9
    # - epochs: 350
    # - learning rate: 0.256, decays by 0.97 every 2.4 epochs
    # - weight decay: 1e-5
    optimizer: torch.optim.Optimizer
    if 'efficientnet' in model_name:
        optimizer = torch.optim.RMSprop(model.parameters(), lr, alpha=0.9,
                                        momentum=0.9, weight_decay=weight_decay)
        lr_scheduler = torch.optim.lr_scheduler.StepLR(
            optimizer=optimizer, step_size=1, gamma=0.97 ** (1 / 2.4))
    else:  # resnet
        optimizer = torch.optim.SGD(model.parameters(), lr, momentum=0.9,
                                    weight_decay=weight_decay)
        lr_scheduler = torch.optim.lr_scheduler.StepLR(
            optimizer=optimizer, step_size=8, gamma=0.1)  # lower every 8 epochs

    best_epoch_metrics: dict[str, float] = {}
    for epoch in range(epochs):
        print(f'Epoch: {epoch}')
        writer.add_scalar('lr', lr_scheduler.get_last_lr()[0], epoch)

        if epoch > 0 and finetune == epoch:
            print('Turning off fine-tune!')
            set_finetune(model, model_name, finetune=False)

        print('- train:')
        train_metrics, train_heaps, train_cm = run_epoch(
            model, loader=loaders['train'], weighted=False, device=device,
            loss_fn=loss_fn, finetune=finetune > epoch, optimizer=optimizer,
            k_extreme=log_extreme_examples)
        train_metrics = prefix_all_keys(train_metrics, prefix='train/')
        log_run('train', epoch, writer, label_names,
                metrics=train_metrics, heaps=train_heaps, cm=train_cm)
        del train_heaps

        print('- val:')
        val_metrics, val_heaps, val_cm = run_epoch(
            model, loader=loaders['val'], weighted=label_weighted,
            device=device, loss_fn=loss_fn, k_extreme=log_extreme_examples)
        val_metrics = prefix_all_keys(val_metrics, prefix='val/')
        log_run('val', epoch, writer, label_names,
                metrics=val_metrics, heaps=val_heaps, cm=val_cm)
        del val_heaps

        lr_scheduler.step()  # decrease the learning rate

        if val_metrics['val/acc_top1'] > best_epoch_metrics.get('val/acc_top1', 0):  # pylint: disable=line-too-long
            filename = os.path.join(logdir, f'ckpt_{epoch}.pt')
            print(f'New best model! Saving checkpoint to {filename}')
            state = {
                'epoch': epoch,
                'model': getattr(model, 'module', model).state_dict(),
                'val/acc': val_metrics['val/acc_top1'],
                'optimizer': optimizer.state_dict()
            }
            torch.save(state, filename)
            best_epoch_metrics.update(train_metrics)
            best_epoch_metrics.update(val_metrics)
            best_epoch_metrics['epoch'] = epoch

            print('- test:')
            test_metrics, test_heaps, test_cm = run_epoch(
                model, loader=loaders['test'], weighted=label_weighted,
                device=device, loss_fn=loss_fn, k_extreme=log_extreme_examples)
            test_metrics = prefix_all_keys(test_metrics, prefix='test/')
            log_run('test', epoch, writer, label_names,
                    metrics=test_metrics, heaps=test_heaps, cm=test_cm)
            del test_heaps

        # stop training after 8 epochs without improvement
        if epoch >= best_epoch_metrics['epoch'] + 8:
            break

    hparams_dict = {
        'model_name': model_name,
        'multilabel': multilabel,
        'finetune': finetune,
        'batch_size': batch_size,
        'epochs': epochs
    }
    metric_dict = prefix_all_keys(best_epoch_metrics, prefix='hparam/')
    writer.add_hparams(hparam_dict=hparams_dict, metric_dict=metric_dict)
    writer.close()

    # do a complete evaluation run
    best_epoch = best_epoch_metrics['epoch']
    evaluate_model.main(
        params_json_path=params_json_path,
        ckpt_path=os.path.join(logdir, f'ckpt_{best_epoch}.pt'),
        output_dir=logdir, splits=evaluate_model.SPLITS)


def log_run(split: str, epoch: int, writer: tensorboard.SummaryWriter,
            label_names: Sequence[str], metrics: MutableMapping[str, float],
            heaps: Mapping[str, Mapping[int, list[HeapItem]]] | None,
            cm: np.ndarray) -> None:
    """Logs the outputs (metrics, confusion matrix, tp/fp/fn images) from a
    single epoch run to Tensorboard.

    Args:
        metrics: dict, keys already prefixed with {split}/
    """
    per_label_recall = recall_from_confusion_matrix(cm, label_names)
    metrics.update(prefix_all_keys(per_label_recall, f'{split}/label_recall/'))

    # log metrics
    for metric, value in metrics.items():
        writer.add_scalar(metric, value, epoch)

    # log confusion matrix
    cm_fig = plot_utils.plot_confusion_matrix(cm, classes=label_names,
                                              normalize=True)
    cm_fig_img = fig_to_img(cm_fig)
    writer.add_image(tag=f'confusion_matrix/{split}', img_tensor=cm_fig_img,
                     global_step=epoch, dataformats='HWC')

    # log tp/fp/fn images
    if heaps is not None:
        for heap_type, heap_dict in heaps.items():
            log_images_with_confidence(writer, heap_dict, label_names,
                                       epoch=epoch, tag=f'{split}/{heap_type}')
    writer.flush()


def log_images_with_confidence(
        writer: tensorboard.SummaryWriter,
        heap_dict: Mapping[int, list[HeapItem]],
        label_names: Sequence[str],
        epoch: int,
        tag: str) -> None:
    """
    Note: performs image normalization in-place

    Args:
        writer: tensorboard.SummaryWriter
        heap_dict: dict, maps label_id to list of HeapItem, where each HeapItem
            data is a tuple (img, target, top3_conf, top3_preds, img_file)
        label_names: list of str, label names in order of label id
        epoch: int
        tag: str
    """
    # for every image: undo normalization, clamp to [0, 1], CHW -> HWC
    # - cannot be in-place, because the HeapItem might be in multiple heaps
    unnormalize = tv.transforms.Normalize(mean=-MEANS/STDS, std=1.0/STDS)
    for label_id, heap in heap_dict.items():
        label_name = label_names[label_id]

        imgs_list = []
        for item in sorted(heap, reverse=True):  # sort largest to smallest
            img = item.data[0].float()  # clamp() only supports fp32 on CPU
            img = unnormalize(img).clamp_(0, 1).permute(1, 2, 0)
            imgs_list.append((img, *item.data[1:]))

        fig, img_files = imgs_with_confidences(imgs_list, label_names)

        # writer.add_figure() has issues => using add_image() instead
        # writer.add_figure(f'{label_name}/{tag}', fig, global_step=epoch)
        writer.add_image(f'{label_name}/{tag}', fig_to_img(fig),
                         global_step=epoch, dataformats='HWC')
        writer.add_text(f'{label_name}/{tag}_files', '\n\n'.join(img_files),
                        global_step=epoch)


def track_extreme_examples(tp_heaps: dict[int, list[HeapItem]],
                           fp_heaps: dict[int, list[HeapItem]],
                           fn_heaps: dict[int, list[HeapItem]],
                           inputs: torch.Tensor,
                           labels: torch.Tensor,
                           img_files: Sequence[str],
                           logits: torch.Tensor,
                           k: int = 5) -> None:
    """Updates the k most extreme true-positive (tp), false-positive (fp), and
    false-negative (fn) examples with examples from this batch.

    Each HeapItem's data attribute is a tuple of:
    - img: torch.Tensor, shape [3, H, W], type float16, values in [0, 1]
    - label: int
    - top3_conf: list of float
    - top3_preds: list of float
    - img_file: str

    Args:
        *_heaps: dict, maps label_id (int) to heap of HeapItems
        inputs: torch.Tensor, shape [batch_size, 3, H, W]
        labels: torch.Tensor, shape [batch_size]
        img_files: list of str
        logits: torch.Tensor, shape [batch_size, num_classes]
        k: int, number of examples to track
    """
    with torch.no_grad():
        inputs = inputs.detach().to(device='cpu', dtype=torch.float16)
        labels_list = labels.tolist()
        batch_probs = torch.nn.functional.softmax(logits, dim=1).cpu()
        zipped = zip(inputs, labels_list, batch_probs, img_files)  # all on CPU
        for img, label, confs, img_file in zipped:
            label_conf = confs[label].item()

            top3_conf, top3_preds = confs.topk(3)
            top3_conf = top3_conf.tolist()
            top3_preds = top3_preds.tolist()

            data = [img, label, top3_conf, top3_preds, img_file]
            if top3_preds[0] == label:  # true positive
                item = HeapItem(priority=label_conf - top3_conf[1], data=data)
                add_to_heap(tp_heaps[label], item, k=k)
            else:
                # false positive for top3_pred[0]
                # false negative for label
                item = HeapItem(priority=top3_conf[0] - label_conf, data=data)
                add_to_heap(fp_heaps[top3_preds[0]], item, k=k)
                add_to_heap(fn_heaps[label], item, k=k)


def correct(outputs: torch.Tensor, labels: torch.Tensor,
            weights: torch.Tensor | None = None,
            top: Sequence[int] = (1,)) -> dict[int, float]:
    """
    Args:
        outputs: torch.Tensor, shape [N, num_classes],
            either logits (pre-softmax) or probabilities
        labels: torch.Tensor, shape [N]
        weights: optional torch.Tensor, shape [N]
        top: tuple of int, list of values of k for calculating top-K accuracy

    Returns: dict, maps k to (weighted) # of correct predictions @ each k
    """
    with torch.no_grad():
        # preds and labels both have shape [N, k]
        _, preds = outputs.topk(k=max(top), dim=1, largest=True, sorted=True)
        labels = labels.view(-1, 1).expand_as(preds)

        corrects = preds.eq(labels).cumsum(dim=1)  # shape [N, k]
        if weights is None:
            corrects = corrects.sum(dim=0)  # shape [k]
        else:
            corrects = weights.matmul(corrects.to(weights.dtype))  # shape [k]
        tops = {k: corrects[k - 1].item() for k in top}
    return tops


def run_epoch(model: torch.nn.Module,
              loader: torch.utils.data.DataLoader,
              weighted: bool,
              device: torch.device,
              top: Sequence[int] = (1, 3),
              loss_fn: torch.nn.Module | None = None,
              finetune: bool = False,
              optimizer: torch.optim.Optimizer | None = None,
              k_extreme: int = 0
              ) -> tuple[
                  dict[str, float],
                  dict[str, dict[int, list[HeapItem]]] | None,
                  np.ndarray
              ]:
    """Runs for 1 epoch.

    Args:
        model: torch.nn.Module
        loader: torch.utils.data.DataLoader
        weighted: bool, whether to use sample weights in calculating loss and
            accuracy
        device: torch.device
        top: tuple of int, list of values of k for calculating top-K accuracy
        loss_fn: optional loss function, calculates per-example loss
        finetune: bool, if true sets model's dropout and BN layers to eval mode
        optimizer: optional optimizer
        k_extreme: int, # of tp/fp/fn examples to track for each label

    Returns:
        metrics: dict, metrics from epoch, contains keys:
            'loss': float, mean per-example loss over entire epoch,
                only included if loss_fn is not None
            'acc_top{k}': float, accuracy@k over the entire epoch
        heaps: dict, keys are ['tp', 'fp', 'fn'], values are heap_dicts,
            each heap_dict maps label_id (int) to a heap of <= 5 HeapItems with
            data attribute (img, target, top3_conf, top3_preds, img_file)
            - 'tp': priority is the difference between target confidence and
                2nd highest confidence
            - 'fp': priority is the difference between highest confidence and
                target confidence
            - 'fn': same as 'fp'
        confusion_matrix: np.ndarray, shape [num_classes, num_classes],
            C[i, j] = # of samples with true label i, predicted as label j
    """
    if optimizer is not None:
        assert loss_fn is not None

    # if evaluating or finetuning, set dropout and BN layers to eval mode
    model.train(optimizer is not None and not finetune)

    if loss_fn is not None:
        losses = AverageMeter()
    accuracies_topk = {k: AverageMeter() for k in top}  # acc@k

    # for each label, track k_extreme most-confident and least-confident images
    if k_extreme > 0:
        tp_heaps: dict[int, list[HeapItem]] = defaultdict(list)
        fp_heaps: dict[int, list[HeapItem]] = defaultdict(list)
        fn_heaps: dict[int, list[HeapItem]] = defaultdict(list)

    all_labels = np.zeros(len(loader.dataset), dtype=np.int32)
    all_preds = np.zeros_like(all_labels)
    end_i = 0

    tqdm_loader = tqdm.tqdm(loader)
    with torch.set_grad_enabled(optimizer is not None):
        for batch in tqdm_loader:
            if weighted:
                inputs, labels, img_files, weights = batch
                weights = weights.to(device, non_blocking=True)
            else:
                # even if batch contains sample weights, don't use them
                inputs, labels, img_files = batch[0:3]
                weights = None

            inputs = inputs.to(device, non_blocking=True)

            batch_size = labels.size(0)
            start_i = end_i
            end_i = start_i + batch_size
            all_labels[start_i:end_i] = labels

            desc = []
            labels = labels.to(device, non_blocking=True)
            outputs = model(inputs)
            all_preds[start_i:end_i] = outputs.detach().argmax(dim=1).cpu()

            if loss_fn is not None:
                loss = loss_fn(outputs, labels)
                if weights is not None:
                    loss *= weights
                loss = loss.mean()
                losses.update(loss.item(), n=batch_size)
                desc.append(f'Loss {losses.val:.4f} ({losses.avg:.4f})')
            if optimizer is not None:
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

            top_correct = correct(outputs, labels, weights=weights, top=top)
            for k, acc in accuracies_topk.items():
                acc.update(top_correct[k] * (100. / batch_size), n=batch_size)
                desc.append(f'Acc@{k} {acc.val:.3f} ({acc.avg:.3f})')
            tqdm_loader.set_description(' '.join(desc))

            if k_extreme > 0:
                track_extreme_examples(tp_heaps, fp_heaps, fn_heaps, inputs,
                                       labels, img_files, outputs, k=k_extreme)

    num_classes = outputs.size(1)
    confusion_matrix = sklearn.metrics.confusion_matrix(
        all_labels, all_preds, labels=np.arange(num_classes))

    metrics = {}
    if loss_fn is not None:
        metrics['loss'] = losses.avg
    for k, acc in accuracies_topk.items():
        metrics[f'acc_top{k}'] = acc.avg
    heaps = None
    if k_extreme > 0:
        heaps = {'tp': tp_heaps, 'fp': fp_heaps, 'fn': fn_heaps}
    return metrics, heaps, confusion_matrix


def _parse_args() -> argparse.Namespace:
    """Parses arguments."""
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description='Trains classifier.')
    parser.add_argument(
        'dataset_dir',
        help='path to directory containing: 1) classification dataset CSV, '
             '2) label index JSON, 3) splits JSON')
    parser.add_argument(
        'cropped_images_dir',
        help='path to local directory where image crops are saved')
    parser.add_argument(
        '--multilabel', action='store_true',
        help='for multi-label, multi-class classification')
    parser.add_argument(
        '-m', '--model-name', default='efficientnet-b0',
        choices=VALID_MODELS,
        help='which EfficientNet or Resnet model')
    parser.add_argument(
        '--pretrained', nargs='?', const=True, default=False,
        help='start with ImageNet pretrained model or a specific checkpoint')
    parser.add_argument(
        '--finetune', type=int, default=0,
        help='only fine tune the final fully-connected layer for the first '
             '<finetune> epochs')
    parser.add_argument(
        '--label-weighted', action='store_true',
        help='weight training samples to balance labels')
    parser.add_argument(
        '--weight-by-detection-conf', nargs='?', const=True, default=False,
        help='weight training examples by detection confidence. '
             'Optionally takes a .npz file for isotonic calibration.')
    parser.add_argument(
        '--epochs', type=int, default=0,
        help='number of epochs for training, 0 for eval-only')
    parser.add_argument(
        '--batch-size', type=int, default=256,
        help='batch size for both training and eval')
    parser.add_argument(
        '--lr', type=float,
        help='initial learning rate, defaults to (0.016 * batch_size / 256)')
    parser.add_argument(
        '--weight-decay', type=float, default=1e-5,
        help='weight decay')
    parser.add_argument(
        '--num-workers', type=int, default=8,
        help='# of workers for data loading')
    parser.add_argument(
        '--logdir', default='.',
        help='directory where TensorBoard logs and a params file are saved')
    parser.add_argument(
        '--log-extreme-examples', type=int, default=0,
        help='# of tp/fp/fn examples to log for each label and split per epoch')
    parser.add_argument(
        '--seed', type=int,
        help='random seed')
    return parser.parse_args()


if __name__ == '__main__':
    args = _parse_args()
    if args.lr is None:
        args.lr = 0.016 * args.batch_size / 256  # based on TF models repo
    main(dataset_dir=args.dataset_dir,
         cropped_images_dir=args.cropped_images_dir,
         multilabel=args.multilabel,
         model_name=args.model_name,
         pretrained=args.pretrained,
         finetune=args.finetune,
         label_weighted=args.label_weighted,
         weight_by_detection_conf=args.weight_by_detection_conf,
         epochs=args.epochs,
         batch_size=args.batch_size,
         lr=args.lr,
         weight_decay=args.weight_decay,
         num_workers=args.num_workers,
         logdir=args.logdir,
         log_extreme_examples=args.log_extreme_examples,
         seed=args.seed)
