import logging
from abc import abstractmethod
from copy import deepcopy
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from uuid import UUID

from highlighter.client.evaluation import EvaluationMetricResult
from highlighter.client.gql_client import HLClient
from highlighter.client.training_runs import TrainingRunType
from highlighter.core.labeled_uuid import LabeledUUID
from highlighter.datasets.dataset import Dataset


class BaseTrainer:
    def __init__(
        self,
        training_run_dir: Path,
        training_run: TrainingRunType,
    ):
        self._logger = logging.getLogger(__name__)
        self._training_run_dir = training_run_dir
        self._training_run = training_run

    @property
    def training_run(self) -> TrainingRunType:
        return self._training_run

    @property
    def _hl_cache_dir(self):
        d = self._training_run_dir / ".hl"
        d.mkdir(exist_ok=True, parents=True)
        return d

    @property
    def _hl_cache_dataset_dir(self):
        d = self._hl_cache_dir / "datasets"
        d.mkdir(exist_ok=True, parents=True)
        return d

    @property
    def training_run_yaml_path(self):
        d = self._training_run_dir / "training_run.yaml"
        return d

    def generate_boilerplate(self, save_training_run_yaml: bool = True):
        """Puts the standard boilerplate files and makes directories required
        by in the Trainer in the training_run_dir. Downloads as caches
        the dataset annotations
        """
        if save_training_run_yaml:
            self._training_run.dump_yaml(self._training_run_dir / "training_run.yaml")

    @property
    @abstractmethod
    def training_data_dir(self) -> Path:
        """Path to data used to store training data"""
        pass

    @abstractmethod
    def generate_trainer_specific_dataset(self, hl_dataset: Dataset):
        pass

    @abstractmethod
    def train(self) -> Any:
        """Trains the Model, return the trained model instance"""
        pass

    @abstractmethod
    def evaluate(
        self, checkpoint: Path, cfg_path: Optional[Path] = None
    ) -> Dict[str, EvaluationMetricResult]:
        """Evaluate a model given a Path to a checkpoint, onnx_file or model instance"""
        pass

    def export(self, checkpoint: Path, cfg_overrides: Optional[Dict] = None) -> List[Path]:
        """Export artefacts by discovering child methods named ``export_*``."""

        if cfg_overrides:
            self._training_run.trainer_config.update(cfg_overrides)

        export_method_names = sorted(name for name in dir(self) if name.startswith("export_"))
        if not export_method_names:
            raise ValueError(f"{self.__class__.__name__} must define at least one export_* method")

        artefact_paths: List[Path] = []
        for method_name in export_method_names:
            export_method = getattr(self, method_name)
            exported = export_method(Path(checkpoint))
            if isinstance(exported, list):
                artefact_paths.extend(Path(p) for p in exported)
            else:
                artefact_paths.append(Path(exported))
        return artefact_paths

    def _resolve_checkpoint_path(self, trained_model: Any) -> Path:
        if isinstance(trained_model, (str, Path)):
            checkpoint_path = Path(trained_model)
            if checkpoint_path.exists():
                return checkpoint_path

        checkpoint_candidates = [
            getattr(getattr(trained_model, "trainer", None), "best", None),
            getattr(getattr(trained_model, "trainer", None), "last", None),
        ]
        for candidate in checkpoint_candidates:
            if candidate is None:
                continue
            candidate_path = Path(candidate)
            if candidate_path.exists():
                return candidate_path

        raise ValueError(f"Unable to determine checkpoint path from {type(trained_model)}")

    @property
    def training_run_id(self) -> int:
        return self._training_run.id

    @property
    def evaluation_id(self) -> int:
        return self._training_run.evaluation_id

    def get_datasets(self, client: HLClient, page_size=20) -> Dataset:
        """Returns a Highlighter SDK Dataset object containing data from each
        split in the one object. `dataset.data_files_id.split` identifies the
        which split each data_file belongs to.
        """
        datasets = Dataset.read_training_run(
            client, self._training_run, self._hl_cache_dataset_dir, page_size=page_size
        )
        return self._combine_hl_datasets(datasets)

    def _combine_hl_datasets(self, datasets):
        # When creating a training run in Highlighter the train split is required
        # but a user can supply either a test or dev set, or both. If not both we
        # duplicate the one that exists here
        if "test" not in datasets:
            datasets["test"] = deepcopy(datasets["dev"])
            datasets["test"].data_files_df.split = "test"
        if "dev" not in datasets:
            datasets["dev"] = deepcopy(datasets["test"])
            datasets["dev"].data_files_df.split = "dev"

        combined_ds = datasets["train"]
        combined_ds.append([datasets["dev"], datasets["test"]])

        return combined_ds

    def filter_dataset(self, dataset: Dataset) -> Dataset:
        """Optionally add some code to filter the Highlighter Datasets as required.
        The YoloWriter will only use entities with both a pixel_location attribute
        and a 'category_attribute_id' attribute when converting to the Yolo dataset format.
        It will the unique values for the object_class attribute as the detection
        categories.

        For example, if you want to train a detector that finds Apples and Bananas,
        and your taxonomy looks like this:

            - object_class: Apple
            - object_class: Orange
            - object_class: Banana

        Then you may do something like this:

            adf = combined_ds.annotations_df
            ddf = combined_ds.data_files_df

            orange_entity_ids = adf[(adf.attribute_id == OBJECT_CLASS_ATTRIBUTE_UUID) &
                                   (adf.value == "Orange")].entity_id.unique()

            # Filter out offending entities
            adf = adf[adf.entity_id.isin(orange_entity_ids)]

            # clean up images that are no longer needed
            ddf = ddf[ddf.data_file_id.isin(adf.data_file_id)]

            combined_ds.annotations_df = adf
        """

        # keep entities with specified categories
        categories = self.get_categories()

        # ToDo: Make sure when values in the dataframes are populated
        # they are a known and consistent primative type.
        def _attr_value(attrs, attr_id):
            if attr_id in attrs:
                return attrs.get(attr_id)
            attr_id_str = str(attr_id)
            if attr_id_str in attrs:
                return attrs.get(attr_id_str)
            try:
                return attrs.get(UUID(attr_id_str))
            except ValueError:
                return None

        def keep_entities_with_categories(e):
            for attr_id, value in categories:
                attr_val = _attr_value(e.attributes, attr_id)
                if attr_val is not None and str(attr_val) == str(value):
                    return True
            return False

        dataset.filter_entities(keep_entities_with_categories)

        return dataset

    def _train(self, hl_dataset: Dataset):
        self.generate_trainer_specific_dataset(hl_dataset)

        trained_model = self.train()
        checkpoint_path = self._resolve_checkpoint_path(trained_model)
        artefact_paths = self.export(checkpoint_path)
        abs_artefact_paths = [artefact_path.absolute() for artefact_path in artefact_paths]
        eval_metric_results = self.evaluate(checkpoint_path)
        return trained_model, abs_artefact_paths, checkpoint_path, eval_metric_results

    def get_categories(self) -> List[Tuple[LabeledUUID, LabeledUUID]]:
        """Return a list of attribute_id, value tuples representing the
        unique categories for training
        """
        head_outputs = self._training_run.input_output_schema.get_head_model_outputs(
            self._training_run.head_idx
        )
        return [
            (
                LabeledUUID(o.taxon[-1].entity_attribute_id, label=o.taxon[-1].entity_attribute_name),
                LabeledUUID(
                    o.taxon[-1].entity_attribute_enum_id, label=o.taxon[-1].entity_attribute_enum_value
                ),
            )
            for o in head_outputs
        ]
