from pydantic import BaseModel, Field
from olmoearth_shared.api.run.evaluation_results import EvaluationMetricType
from uuid import UUID
CHECKPOINT_FILE_NAME = "checkpoint.ckpt"
TRAINER_LAST_CHECKPOINT_FILE_NAME = "last.ckpt"
MODEL_CONFIG_FILE_NAME = "model.yaml"
DATASET_CONFIG_FILE_NAME = "dataset.json"
DATASET_CONFIG_FILE_NAME_IN_DATASET_DIR = "config.json"
OLMOEARTH_RUN_CONFIG_FILE_NAME = "olmoearth_run.yaml"
ANNOTATION_FEATURES_FILE_NAME = "annotation_features.geojson"
ANNOTATION_TASK_FEATURES_FILE_NAME = "annotation_task_features.geojson"
EXTRA_MODEL_FILES_DIR_NAME = "extra_model_files"
TRAINER_CHECKPOINTS_DIR_NAME = "trainer_checkpoints"
MODEL_EVALUATION_METADATA_FILE_NAME = "model_evaluation_metadata.json"

EVALUATION_DIR_NAME = "evaluation"
CONFUSION_MATRIX_EXAMPLES_DIR = "confusion_matrix_examples"

class ModelStagePaths(BaseModel):
    root_path: str = Field(description="A path to the location containing the configuration files that define the model")

    @property
    def default_dataset_path(self) -> str:
        """Path to the dataset for training/inference"""
        return f"{self.root_path}/dataset"

    @property
    def checkpoint_path(self) -> str:
        """Path to the checkpoint file containing the trained weights and biases"""
        return f"{self.root_path}/{CHECKPOINT_FILE_NAME}"

    @property
    def trainer_checkpoints_path(self) -> str:
        """
        Path lightning manages its in-flight state in during training, including:
        intermediate checkpoint files, trainer state, etc.
        """
        return f"{self.root_path}/{TRAINER_CHECKPOINTS_DIR_NAME}"

    @property
    def trainer_last_checkpoint_path(self) -> str:
        """Path to the last checkpoint file saved by the trainer"""
        return f"{self.trainer_checkpoints_path}/{TRAINER_LAST_CHECKPOINT_FILE_NAME}"

    @property
    def model_config_path(self) -> str:
        """Path to the model config that was used to train this ModelStage"""
        return f"{self.root_path}/{MODEL_CONFIG_FILE_NAME}"

    @property
    def dataset_config_path(self) -> str:
        """Path to the dataset config that defines the windows of this ModelStage"""
        return f"{self.root_path}/{DATASET_CONFIG_FILE_NAME}"

    @property
    def dataset_config_path_in_dataset_dir(self) -> str:
        """The configuration file where it needs to be when actually consumed by rslearn"""
        return f"{self.default_dataset_path}/{DATASET_CONFIG_FILE_NAME_IN_DATASET_DIR}"

    @property
    def olmoearth_run_config_path(self) -> str:
        """
        Path to the OlmoEarthRun config that defines:
            1. Partition strategies: How to split the Workflow into partitioned Tasks; and within a Task how to build individual window `STGeometry`s from a partition
            2. Postprocessing strategies: How to postprocess individual Windows, combine Windows within a Partition, and combine Partitions
            3. Fine tuning window preparation
        """
        return f"{self.root_path}/{OLMOEARTH_RUN_CONFIG_FILE_NAME}"

    @property
    def annotation_features_path(self) -> str:
        """Path to the annotation features GeoJSON file for fine-tuning"""
        return f"{self.root_path}/{ANNOTATION_FEATURES_FILE_NAME}"

    @property
    def annotation_task_features_path(self) -> str:
        """Path to the annotation task features GeoJSON file for fine-tuning"""
        return f"{self.root_path}/{ANNOTATION_TASK_FEATURES_FILE_NAME}"

    @property
    def extra_model_files_path(self) -> str:
        """
        Path to any extra files required to train or operate this model,
        e.g. pretrained embedding model weights + config, preprocessing config
        """
        return f"{self.root_path}/{EXTRA_MODEL_FILES_DIR_NAME}"

    @property
    def model_evaluation_metadata_path(self) -> str:
        """Path to the metadata JSON saved during model evaluation prediction."""
        return f"{self.default_dataset_path}/{MODEL_EVALUATION_METADATA_FILE_NAME}"

    @property
    def evaluation_dir_path(self) -> str:
        """Path to the evaluation directory for this model stage"""
        return f"{self.root_path}/{EVALUATION_DIR_NAME}"

    def get_evaluation_metric_path(self, metric_type: EvaluationMetricType) -> str:
        """
        Path to a specific evaluation metric file.

        IMPORTANT: This method assumes a 1:1 mapping between metric_type and filename:
        - metric_type="confusion_matrix" -> confusion_matrix.json
        - metric_type="iou" -> iou.json

        This convention must be followed when:
        1. Adding new metric types to EvaluationMetricType
        2. Writing evaluation results to GCS
        3. Reading evaluation results in API routes

        """
        return f"{self.evaluation_dir_path}/{metric_type}.json"

    @property
    def confusion_matrix_path(self) -> str:
        """Path to the confusion matrix evaluation file with embedded provenance"""
        return self.get_evaluation_metric_path(EvaluationMetricType.CONFUSION_MATRIX)

    @property
    def confusion_matrix_examples_dir_path(self) -> str:
        """Path to the directory containing per-cell confusion matrix example files"""
        return f"{self.evaluation_dir_path}/{CONFUSION_MATRIX_EXAMPLES_DIR}"

    def get_confusion_matrix_cell_examples_path(self, predicted_class_id: UUID, actual_class_id: UUID) -> str:
        """
        Path to a specific confusion matrix cell's examples file.

        Args:
            predicted_class_id: The predicted class ID
            actual_class_id: The actual/ground truth class ID

        Returns:
            Path like: {root}/evaluation/confusion_matrix_examples/predicted_0_actual_1.json
        """
        examples_path = f"{self.confusion_matrix_examples_dir_path}"
        return f"{examples_path}/predicted_{predicted_class_id}_actual_{actual_class_id}.json"

    def get_evaluation_metric_examples_path(self, metric_type: EvaluationMetricType, predicted_class_id: UUID, actual_class_id: UUID) -> str:
        """
        Path to a specific evaluation metric examples file.
        """
        if metric_type == EvaluationMetricType.CONFUSION_MATRIX:
            return self.get_confusion_matrix_cell_examples_path(predicted_class_id, actual_class_id)
        else:
            raise ValueError(f"Unsupported metric type: {metric_type}")
