"""Read and query calendar-dated disease-event arrays from prepare_event_dates.py."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

import numpy as np
import pandas as pd


REQUIRED_FIELDS = {"eid", "event_date", "token"}


def load_event_dates(path: str | Path) -> np.ndarray:
    """Load and validate the structured ``.npy`` event array."""
    events = np.load(path)
    if events.dtype.names is None or not REQUIRED_FIELDS.issubset(events.dtype.names):
        raise ValueError(
            "Expected a structured .npy with eid, event_date, token fields. "
            "Create it with prepare_event_dates.py."
        )
    return events


def load_token_labels(labels_file: str | Path) -> dict[int, str]:
    """Load token -> human-readable code using the project label convention."""
    labels = {1: "CHECKUP"}
    with Path(labels_file).open(encoding="utf-8") as handle:
        for index, line in enumerate(handle):
            code = line.strip().split(" ", maxsplit=1)[0]
            if code:
                labels[index + 2] = code
    return labels


@dataclass
class EventDateIndex:
    """Small in-memory query wrapper for exposure-linkage and cohort scripts."""

    events: np.ndarray
    token_labels: dict[int, str] | None = None

    @classmethod
    def from_files(
        cls,
        event_file: str | Path,
        labels_file: str | Path | None = None,
    ) -> "EventDateIndex":
        labels = load_token_labels(labels_file) if labels_file is not None else None
        return cls(load_event_dates(event_file), labels)

    def to_frame(self, events: np.ndarray | None = None) -> pd.DataFrame:
        """Convert records to a convenient, calendar-dated DataFrame."""
        data = self.events if events is None else events
        frame = pd.DataFrame(
            {
                "eid": data["eid"].astype("int64"),
                "event_date": pd.to_datetime(data["event_date"]),
                "token": data["token"].astype("int32"),
            }
        )
        if self.token_labels is not None:
            frame["label_code"] = frame["token"].map(self.token_labels).fillna("UNKNOWN")
        return frame.sort_values(["eid", "event_date", "token"], kind="stable").reset_index(drop=True)

    def for_eid(self, eid: int) -> pd.DataFrame:
        """Return every stored disease/death event for one participant."""
        return self.to_frame(self.events[self.events["eid"] == int(eid)])

    def between(
        self,
        start: str | pd.Timestamp,
        end: str | pd.Timestamp,
        *,
        eids: Iterable[int] | None = None,
        tokens: Iterable[int] | None = None,
    ) -> pd.DataFrame:
        """Query events in an inclusive calendar-date interval."""
        start_day = np.datetime64(pd.Timestamp(start).date(), "D")
        end_day = np.datetime64(pd.Timestamp(end).date(), "D")
        mask = (self.events["event_date"] >= start_day) & (self.events["event_date"] <= end_day)
        if eids is not None:
            mask &= np.isin(self.events["eid"], list(eids))
        if tokens is not None:
            mask &= np.isin(self.events["token"], list(tokens))
        return self.to_frame(self.events[mask])

    def anchors_before(self, eid: int, date: str | pd.Timestamp) -> pd.DataFrame:
        """Return a participant's event history strictly before an exposure anchor."""
        day = np.datetime64(pd.Timestamp(date).date(), "D")
        mask = (self.events["eid"] == int(eid)) & (self.events["event_date"] < day)
        return self.to_frame(self.events[mask])

    def first_event(self, token: int) -> pd.DataFrame:
        """Return each participant's first date for a requested token."""
        selected = self.events[self.events["token"] == int(token)]
        # Arrays produced by prepare_event_dates.py are already deduplicated;
        # sorting makes this safe for externally produced compatible arrays too.
        order = np.lexsort((selected["event_date"], selected["eid"]))
        selected = selected[order]
        _, first = np.unique(selected["eid"], return_index=True)
        return self.to_frame(selected[first])