"""
Challenge Dataset Loader for the SNOMED CT Entity Linking Challenge.

Provides programmatic access to official challenge notes and annotations:
- snomed-ct-entity-linking-challenge-1.2.0/train_notes.csv
- snomed-ct-entity-linking-challenge-1.2.0/train_annotations.csv
"""

import os
from typing import Dict, List, Optional, Tuple, Any
import pandas as pd

CHALLENGE_DIR = os.path.join(
    os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
    "snomed-ct-entity-linking-challenge-1.2.0",
)
NOTES_FILE = os.path.join(CHALLENGE_DIR, "train_notes.csv")
ANNOTATIONS_FILE = os.path.join(CHALLENGE_DIR, "train_annotations.csv")


class ChallengeDatasetLoader:
    """Loads and inspects clinical notes and annotations from the SNOMED challenge dataset."""

    def __init__(self, notes_path: str = NOTES_FILE, annotations_path: str = ANNOTATIONS_FILE):
        self.notes_path = notes_path
        self.annotations_path = annotations_path
        self._df_notes: Optional[pd.DataFrame] = None
        self._df_annos: Optional[pd.DataFrame] = None

    @property
    def notes(self) -> pd.DataFrame:
        if self._df_notes is None:
            if not os.path.exists(self.notes_path):
                raise FileNotFoundError(f"Challenge notes file not found: {self.notes_path}")
            self._df_notes = pd.read_csv(self.notes_path)
        return self._df_notes

    @property
    def annotations(self) -> pd.DataFrame:
        if self._df_annos is None:
            if not os.path.exists(self.annotations_path):
                raise FileNotFoundError(f"Challenge annotations file not found: {self.annotations_path}")
            self._df_annos = pd.read_csv(self.annotations_path)
        return self._df_annos

    def list_note_ids(self) -> List[str]:
        """Returns list of all note IDs present in the challenge dataset."""
        return self.notes["note_id"].tolist()

    def get_note_by_id(self, note_id: str) -> Optional[str]:
        """Returns the full text of a specific clinical note by note_id."""
        sub = self.notes[self.notes["note_id"] == note_id]
        if sub.empty:
            return None
        return str(sub.iloc[0]["text"])

    def get_annotations_by_note_id(self, note_id: str) -> pd.DataFrame:
        """Returns all ground-truth SNOMED annotations for a specific note."""
        return self.annotations[self.annotations["note_id"] == note_id].copy()

    def get_note_summary(self, note_id: str) -> Dict[str, Any]:
        """Returns note metadata, text preview, and ground truth annotations summary."""
        text = self.get_note_by_id(note_id)
        if text is None:
            raise ValueError(f"Note '{note_id}' not found in dataset.")
        annos = self.get_annotations_by_note_id(note_id)
        return {
            "note_id": note_id,
            "char_length": len(text),
            "total_annotations": len(annos),
            "unique_concepts": annos["concept_id"].nunique() if not annos.empty else 0,
            "text": text,
            "annotations": annos.to_dict(orient="records"),
        }


def load_challenge_note(note_id: str = "10043750-DS-6") -> Dict[str, Any]:
    """Convenience helper to retrieve a note and its annotations."""
    loader = ChallengeDatasetLoader()
    return loader.get_note_summary(note_id)