import json
from collections import defaultdict
from pathlib import Path
from typing import Any, Dict, List, Optional
from uuid import UUID

import numpy as np
import ultralytics
import yaml
from ultralytics import cfg as ultralytics_cfg

from highlighter.client import TrainingRunType, json_tools
from highlighter.client.evaluation import (
    EvaluationMetric,
    EvaluationMetricCodeEnum,
    EvaluationMetricResult,
    find_or_create_evaluation_metric,
)
from highlighter.client.gql_client import HLClient
from highlighter.client.io import multithread_graphql_file_download
from highlighter.datasets.cropping import CropArgs
from highlighter.datasets.formats.yolo.writer import YoloWriter
from highlighter.trainers.base_trainer import BaseTrainer

__all__ = ["YoloV11Trainer"]


class EvaluationCharts:
    PER_CLASS = "Per Class Metrics"
    AGG = "Aggregate Metrics"
    MODEL_SIZE = "Model Size"


class YoloV11Trainer(BaseTrainer):
    def __init__(
        self,
        training_run_dir: Path,
        training_run: TrainingRunType,
    ):
        super().__init__(training_run_dir, training_run)

    def _combine_hl_datasets(self, datasets):
        combined_ds = super()._combine_hl_datasets(datasets)

        # Ultralytics name their dataset splits differently so
        # we need to map them
        #   their "val" is our "test"
        #   their "test" is our "dev"
        combined_ds.data_files_df.loc[combined_ds.data_files_df.split == "test", "split"] = "val"
        combined_ds.data_files_df.loc[combined_ds.data_files_df.split == "dev", "split"] = "test"
        return combined_ds

    def _get_hl_metric(
        self, client, code, chart, desc, name, iou: Optional[float] = None, cat_uuid: Optional[UUID] = None
    ):
        metric = EvaluationMetric(
            evaluation_id=self.evaluation_id,
            code=code,
            chart=chart,
            description=desc,
            name=name,
            object_class_uuid=cat_uuid,
            iou=iou,
        )
        return find_or_create_evaluation_metric(client, metric)[0]

    def _make_classify_onnx_artefact(self, onnx_filepath: Path) -> Path:
        if isinstance(self._training_run.crop_args, dict):
            crop_args_dict = self._training_run.crop_args
        elif self._training_run.crop_args is not None:
            crop_args_dict = self._training_run.crop_args.model_dump()
        else:
            crop_args_dict = None

        d = dict(
            file_url=str(Path(onnx_filepath.absolute())),
            type="OnnxOpset14",
            inference_config=dict(
                type="classifier",
                code="BoxClassifier",
                machine_agent_type_id="d4787671-3839-4af9-9b34-a686faafbfae",
                parameters=dict(
                    output_format="yolov8_cls",
                    cropper=crop_args_dict,
                ),
            ),
        )

        artefact_path = onnx_filepath.parent / "artefact.onnx.yaml"
        with artefact_path.open("w") as f:
            yaml.dump(d, f)

        return artefact_path

    def _make_detect_segment_onnx_artefact(self, onnx_filepath: Path) -> Path:
        output_format = (
            "yolov8_seg" if self._training_run.trainer_config["task"] == "segment" else "yolov8_det"
        )
        d = dict(
            file_url=str(Path(onnx_filepath.absolute())),
            type="OnnxOpset14",
            inference_config=dict(
                type="detector",
                code="Detector",
                machine_agent_type_id="29653174-8f45-440d-b75a-4ed0aa5fa6ff",
                parameters=dict(
                    output_format=output_format,
                ),
            ),
        )

        artefact_path = onnx_filepath.parent / "artefact.onnx.yaml"
        with artefact_path.open("w") as f:
            yaml.dump(d, f)

        return artefact_path

    def export_onnx(self, checkpoint: Path) -> Path:
        # Disable ultralyitcs' auto install of packages
        ultralytics.utils.checks.AUTOINSTALL = False

        trained_model = ultralytics.YOLO(str(checkpoint), task=self._training_run.trainer_config["task"])
        onnx_filepath = Path(trained_model.export(format="onnx", dynamic=True, device=0, batch=1))

        task = self._training_run.trainer_config["task"]
        if task == "classify":
            return self._make_classify_onnx_artefact(onnx_filepath)
        if task in ("detect", "segment"):
            return self._make_detect_segment_onnx_artefact(onnx_filepath)
        raise ValueError(f"Invalid yolo task '{task}', expected one of (classify|detect|segment)")

    def export_pytorch(self, checkpoint: Path) -> Path:
        task = self._training_run.trainer_config["task"]

        if task == "classify":
            if isinstance(self._training_run.crop_args, dict):
                crop_args_dict = self._training_run.crop_args
            elif self._training_run.crop_args is not None:
                crop_args_dict = self._training_run.crop_args.model_dump()
            else:
                crop_args_dict = None

            inference_config = dict(
                type="classifier",
                code="BoxClassifier",
                machine_agent_type_id="d4787671-3839-4af9-9b34-a686faafbfae",
                parameters=dict(
                    output_format="yolov8_cls",
                    cropper=crop_args_dict,
                ),
            )
        elif task in ("detect", "segment"):
            output_format = "yolov8_seg" if task == "segment" else "yolov8_det"
            inference_config = dict(
                type="detector",
                code="Detector",
                machine_agent_type_id="29653174-8f45-440d-b75a-4ed0aa5fa6ff",
                parameters=dict(
                    output_format=output_format,
                ),
            )
        else:
            raise ValueError(f"Invalid yolo task '{task}', expected one of (classify|detect|segment)")

        d = dict(
            file_url=str(checkpoint.absolute()),
            type="PyTorch",
            inference_config=inference_config,
        )

        artefact_path = checkpoint.parent / "artefact.pytorch.yaml"
        with artefact_path.open("w") as f:
            yaml.dump(d, f)

        return artefact_path

    def train(self) -> Any:
        model = ultralytics.YOLO(self._training_run.trainer_config["model"])

        ultralytics.settings.update({"datasets_dir": str(self._training_run_dir.absolute())})

        data_cfg_path = (self._training_run_dir / "datasets" / "data.yaml").absolute()
        with data_cfg_path.open("r") as f:
            data_cfg = yaml.safe_load(f)

        self._training_run.trainer_config["data"] = str(data_cfg_path)
        self._training_run.trainer_config["single_cls"] = data_cfg["nc"] == 1

        if self._training_run.trainer_config["task"] == "classify":
            self._training_run.trainer_config["data"] = str(data_cfg_path.parent)
            self._training_run.trainer_config["classes"] = (
                None  # Classification classes are specified via directories
            )
        else:
            self._training_run.trainer_config["classes"] = list(data_cfg["names"].keys())
            self._training_run.trainer_config["data"] = str(data_cfg_path)

        trainer = None

        if getattr(self._training_run, "oversample", True):
            task = self._training_run.trainer_config.get("task")

            if task == "classify":
                from highlighter.trainers.yolov11._weighted_sampling import (
                    WeightedClassificationTrainer,
                )

                trainer = WeightedClassificationTrainer

            elif task == "detect" or task == "segment":
                from highlighter.trainers.yolov11._weighted_sampling import (
                    WeightedDetectionTrainer,
                )

                trainer = WeightedDetectionTrainer

            else:
                raise ValueError(
                    f"Oversampling is enabled but task '{task}' is not supported. "
                    "Supported tasks: ['classify', 'detect', 'segment']"
                )

        model.train(trainer=trainer, **self._training_run.trainer_config)
        return model

    @property
    def training_data_dir(self) -> Path:
        return self._training_run_dir / "datasets"

    def generate_trainer_specific_dataset(self, hl_dataset):
        if not (self.training_data_dir / "data.yaml").exists():
            self.training_data_dir.mkdir(exist_ok=True)
            # Optionally filter dataset, see filter_dataset's doc str
            filtered_hl_ds = self.filter_dataset(hl_dataset)

            # Download required images
            image_cache_dir = self._training_run_dir / "images"
            multithread_graphql_file_download(
                HLClient.get_client(),
                filtered_hl_ds.data_files_df.data_file_id.values,
                image_cache_dir,
            )

            ddf = filtered_hl_ds.data_files_df
            if any([Path(f).suffix.lower() == ".mp4" for f in ddf.filename.unique()]):
                print("Detected video dataset, interpolating data from keyframes")
                filtered_hl_ds = filtered_hl_ds.interpolate_from_key_frames(
                    frame_save_dir=image_cache_dir,
                    source_file_dir=image_cache_dir,
                )

            # Write dataset in yolo format
            writer = YoloWriter(
                output_dir=self.training_data_dir,
                image_cache_dir=image_cache_dir,
                categories=self.get_categories(),
                task=self._training_run.trainer_config["task"],
                crop_args=self._training_run.crop_args,
            )
            writer.write(filtered_hl_ds)

    def evaluate(
        self, checkpoint: Path | str, cfg_path: Optional[Path | str] = None
    ) -> Dict[str, EvaluationMetricResult]:
        """Evaluate a YOLOv11 model and generate structured evaluation metrics.

        Evaluates the model using ultralytics validation and creates evaluation metric
        result objects for upload to Highlighter. Handles both classification and
        detection/segmentation tasks with appropriate metrics for each.

        Args:
            checkpoint: Path to model checkpoint (can be .pt, .onnx, or other supported formats)
            cfg_path: Optional path to configuration file (uses self._cfg if None)

        Returns:
            List[EvaluationMetricResult]: Evaluation metric results ready for upload

        Side Effects:
            Saves evaluation results to 'eval_metric_results.json' in checkpoint directory
        """
        eval_metric_results: Dict[str, EvaluationMetricResult] = {}
        if cfg_path is None:
            cfg: Dict = self._training_run.trainer_config
        else:
            cfg: Dict = ultralytics.utils.YAML.load(str(cfg_path))

        assert cfg["data"] is not None
        client = HLClient.get_client()

        model = ultralytics.YOLO(str(checkpoint), task=cfg["task"])
        results = model.val(**cfg)

        if cfg["task"] in ("segment", "detect"):
            eval_metric_results = self._create_eval_results_det_or_seg(results, cfg, client)
        elif cfg["task"] == "classify":
            eval_metric_results = self._create_eval_results_cls(results, cfg, client)

        with (Path(checkpoint).parent / "eval_metric_results.json").open("w") as f:
            json.dump(
                {k: e.model_dump() for k, e in eval_metric_results.items()},
                f,
                indent=2,
                cls=json_tools.HLJSONEncoder,
            )

        return eval_metric_results

    def _create_eval_results_cls(
        self, results, cfg: Dict, client: HLClient
    ) -> Dict[str, EvaluationMetricResult]:
        """Create evaluation metric results for classification tasks.

        Extracts classification-specific metrics including model size, top-1 accuracy,
        and top-5 accuracy from ultralytics validation results.

        Args:
            results: Ultralytics validation results object
            cfg: Configuration dictionary containing model information
            client: HLClient for metric lookup

        Returns:
            List[EvaluationMetricResult]: Classification evaluation metrics
        """
        eval_metric_results: Dict[str, EvaluationMetricResult] = {}

        model_size_char = Path(cfg["model"]).stem.replace("-cls", "")[-1]
        if model_size_char in "nsmlx":
            model_size_int = "nsmlx".index(model_size_char)

            name = "Model Size"
            eval_metric_results[name] = self._get_hl_metric(
                client=client,
                code=EvaluationMetricCodeEnum.Other,
                chart=EvaluationCharts.AGG,
                desc="Size of the model, 0:Nano|1:Small|2:Medium|3:Large|4:XLarge",
                name=name,
            ).result(model_size_int, self.training_run_id)
        else:
            model_str = cfg["model"]
            raise SystemExit(f"Unvalid model_size_char '{model_size_char}' from '{model_str}'")

        name = "Accuracy Top1"
        acc_top1 = results.results_dict["metrics/accuracy_top1"]
        eval_metric_results[name] = self._get_hl_metric(
            client=client,
            code=EvaluationMetricCodeEnum.Other,
            chart=EvaluationCharts.AGG,
            desc="Top One Accuracy",
            name=name,
        ).result(acc_top1, self.training_run_id)

        acc_top5 = results.results_dict["metrics/accuracy_top5"]
        eval_metric_results[name] = self._get_hl_metric(
            client=client,
            code=EvaluationMetricCodeEnum.Other,
            chart=EvaluationCharts.AGG,
            desc="Top Five Accuracy",
            name=name,
        ).result(acc_top5, self.training_run_id)

        per_class_P, per_class_R = self._precision_recall_from_cm(
            results.confusion_matrix.matrix,
        )

        for idx, cat in enumerate(self.get_categories()):
            name = f"Precision({cat.label})"
            eval_metric_results[name] = self._get_hl_metric(
                client=client,
                code=EvaluationMetricCodeEnum.Other,
                chart=EvaluationCharts.PER_CLASS,
                desc=f"One vs All Precision of ({cat.short_str()})",
                name=name,
                cat_uuid=cat,
            ).result(per_class_P[idx], self.training_run_id)

            name = f"Recall({cat.label})"
            eval_metric_results[name] = self._get_hl_metric(
                client=client,
                code=EvaluationMetricCodeEnum.Other,
                chart=EvaluationCharts.PER_CLASS,
                desc=f"One vs All Recall of ({cat.short_str()})",
                name=name,
                cat_uuid=cat,
            ).result(per_class_R[idx], self.training_run_id)

        name = "Average Precision"
        ave_P = sum(per_class_P.values()) / len(per_class_P)
        eval_metric_results[name] = self._get_hl_metric(
            client=client,
            code=EvaluationMetricCodeEnum.Other,
            chart=EvaluationCharts.AGG,
            desc="Average Precision",
            name=name,
        ).result(ave_P, self.training_run_id)

        return eval_metric_results

    def _precision_recall_from_cm(self, cm):
        if cm.shape != (3, 3):
            raise ValueError("Confusion matrix must be 3x3")

        precision = {}
        recall = {}

        for i, cat in enumerate(self.get_categories()):
            TP = cm[i, i]
            FP = np.sum(cm[:, i]) - TP  # sum of column i, minus TP
            FN = np.sum(cm[i, :]) - TP  # sum of row i, minus TP

            # Avoid division by zero
            precision[i] = TP / (TP + FP) if (TP + FP) > 0 else 0.0
            recall[i] = TP / (TP + FN) if (TP + FN) > 0 else 0.0

        return precision, recall

    def _create_eval_results_det_or_seg(
        self, results, cfg: Dict, client: HLClient
    ) -> Dict[str, EvaluationMetricResult]:
        """Create evaluation metric results for detection and segmentation tasks.

        Extracts detection/segmentation-specific metrics including precision, recall,
        mAP@IoU50, model size, ideal confidence threshold, and per-class mAP values
        from ultralytics validation results.

        Args:
            results: Ultralytics validation results object
            cfg: Configuration dictionary containing model information
            client: HLClient for metric lookup

        Returns:
            List[EvaluationMetricResult]: Detection/segmentation evaluation metrics

        Note:
            Uses corrected ultralytics result attributes (e.g., 'metrics/precision(B)')
            rather than deprecated mean_results() method for better accuracy.
        """
        eval_metric_results: Dict[str, EvaluationMetricResult] = {}

        model_size_char = Path(cfg["model"]).stem.replace("-cls", "")[-1]
        if model_size_char in "nsmlx":
            model_size_int = "nsmlx".index(model_size_char)

            name = "Model Size"
            eval_metric_results[name] = self._get_hl_metric(
                client=client,
                code=EvaluationMetricCodeEnum.Other,
                chart=EvaluationCharts.AGG,
                desc="Size of the model, 0:Nano|1:Small|2:Medium|3:Large|4:XLarge",
                name=name,
            ).result(model_size_int, self.training_run_id)
        else:
            model_str = cfg["model"]
            raise SystemExit(f"Unvalid model_size_char '{model_size_char}' from '{model_str}'")

        name = "Precision"
        eval_metric_results[name] = self._get_hl_metric(
            client,
            code=EvaluationMetricCodeEnum.Other,
            chart=EvaluationCharts.AGG,
            desc="Precision Over All Classes",
            name=name,
            iou=0.5,
        ).result(
            results.results_dict["metrics/precision(B)"],
            self.training_run_id,
        )

        name = "Recall"
        eval_metric_results[name] = self._get_hl_metric(
            client,
            code=EvaluationMetricCodeEnum.Other,
            chart=EvaluationCharts.AGG,
            desc="Recall Over All Classes",
            name=name,
            iou=0.5,
        ).result(results.results_dict["metrics/recall(B)"], self.training_run_id)

        name = "mAP@IOU50"
        eval_metric_results[name] = self._get_hl_metric(
            client,
            code=EvaluationMetricCodeEnum.mAP,
            chart=EvaluationCharts.AGG,
            desc="Mean Avarage Precision Over All Classes",
            iou=0.5,
            name=name,
        ).result(results.results_dict["metrics/mAP50(B)"], self.training_run_id)

        f1_conf_curve_idx = results.curves.index("F1-Confidence(B)")

        if len(results.curves_results[f1_conf_curve_idx][1].shape) == 2:
            best_f1_idx = np.argmax(results.curves_results[f1_conf_curve_idx][1].mean(axis=0))
        else:
            best_f1_idx = np.argmax(results.curves_results[f1_conf_curve_idx][1])
        best_f1_thr = results.curves_results[f1_conf_curve_idx][0][best_f1_idx]

        name = "Ideal Confidence Threshold"
        eval_metric_results[name] = self._get_hl_metric(
            client,
            code=EvaluationMetricCodeEnum.Other,
            chart="Ideal Confidence Threshold",
            desc="The Confidence Threshold that maximizes the F1Score",
            name=name,
        ).result(best_f1_thr, self.training_run_id)

        cats = [c[1] for c in self.get_categories()]
        CLS_PRECISION_IOU50 = 2
        for cls_idx, cat in enumerate(cats):
            class_name = cat.label
            name = f"mAP@IOU50({class_name})"
            eval_metric_results[name] = self._get_hl_metric(
                client,
                code=EvaluationMetricCodeEnum.mAP,
                chart=EvaluationCharts.AGG,
                desc=f"Mean Avarage Precision ({cat.short_str()})",
                name=name,
                iou=0.5,
                cat_uuid=cat,
            ).result(results.class_result(cls_idx)[CLS_PRECISION_IOU50], self.training_run_id)

        return eval_metric_results
