"""
Classification validator: accuracy, top-1/top-5, confusion matrix.
Uses Ultralytics ClassifyMetrics and ConfusionMatrix.
"""

from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import numpy as np
import torch

from .dataloader import DatasetLoader
from .metrics import ClassifyMetrics, ConfusionMatrix


class ClassifyValidator:
    def __init__(
        self,
        model: Any,
        val_txt_path: str,
        batch_size: int = 1,
        save_dir: Union[str, Path] = ".",
        plots: bool = False,
        names: Optional[Dict[int, str]] = None,
        nc: int = 1000,
        img_size: tuple = (224, 224),
        verbose: bool = True,
    ):
        self.model = model
        self.val_txt_path = val_txt_path
        self.batch_size = batch_size
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(parents=True, exist_ok=True)
        self.plots = plots
        self.names = names or {}
        self.nc = nc
        self.img_size = img_size
        self.verbose = verbose
        self.metrics = ClassifyMetrics(names=self.names)
        self.confusion_matrix = ConfusionMatrix(names=self.names, task="classify", save_matches=False)
        self.pred: List[torch.Tensor] = []
        self.targets: List[torch.Tensor] = []

    def _get_loader(self) -> DatasetLoader:
        return DatasetLoader(
            self.val_txt_path,
            task="classify",
            num_kpts=None,
            batch_size=self.batch_size,
            img_size=self.img_size,
        )

    def __call__(self) -> Dict[str, Any]:
        loader = self._get_loader()
        self.pred = []
        self.targets = []
        for batch_imgs, batch_gt, batch_paths in loader:
            preds = self.model(batch_imgs)
            # preds: list of per-image outputs; classify is (num_classes,) or (1, num_classes)
            for si in range(len(batch_imgs)):
                if isinstance(preds, list):
                    p = preds[si]
                else:
                    p = preds[si] if preds.ndim > 1 else preds
                if isinstance(p, np.ndarray):
                    p = torch.from_numpy(p)
                if p.ndim == 1:
                    p = p.unsqueeze(0)
                n5 = min(p.shape[1], 5)
                top5 = p.argsort(dim=1, descending=True)[:, :n5].type(torch.int32).cpu()
                self.pred.append(top5)
                gt = batch_gt[si]
                if isinstance(gt, np.ndarray):
                    gt = torch.from_numpy(gt)
                if gt.dim() == 0:
                    gt = gt.unsqueeze(0)
                self.targets.append(gt.type(torch.int32).cpu())

        self.metrics.process(self.targets, self.pred)
        self.confusion_matrix.process_cls_preds(self.pred, self.targets)
        if self.plots:
            for normalize in (True, False):
                self.confusion_matrix.plot(save_dir=str(self.save_dir), normalize=normalize)
        self.metrics.confusion_matrix = self.confusion_matrix
        self.metrics.save_dir = self.save_dir
        if self.verbose:
            print(self.metrics.results_dict)
        return self.metrics.results_dict
