ntml_parser.py

from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import Iterable, List, Optional, Set

TIMEFRAME_PATTERN = re.compile(r"\b(\d{4}-\d{4})\b")
MIT_PATTERN = re.compile(r"(\d+)\s*MIT", re.IGNORECASE)
MINIT_PATTERN = re.compile(r"(\d+)\s*MINIT", re.IGNORECASE)
PROVIDING_PATTERN = re.compile(r"^[A-Z]{3,4}$")


@dataclass
class Restriction:
    raw: str
    requesting_facility: Optional[str] = None
    providing_facilities: List[str] = field(default_factory=list)
    timeframe: Optional[str] = None
    mit_value: Optional[int] = None
    minit_value: Optional[int] = None
    impact_conditions: List[str] = field(default_factory=list)


@dataclass
class ParsedRestrictions:
    restrictions: List[Restriction]
    min_mit: Optional[int]
    max_mit: Optional[int]
    min_minit: Optional[int]
    max_minit: Optional[int]
    requesting_facilities: Set[str]
    providing_facilities: Set[str]
    impact_conditions: Set[str]
    timeframes: Set[str]


def _extract_requesting_facility(line: str) -> Optional[str]:
    if " via " in line:
        return line.split(" via ", 1)[0].strip() or None
    return None


def _parse_providing_facilities(tokens: Iterable[str]) -> List[str]:
    facilities: List[str] = []
    for token in tokens:
        if ":" not in token:
            continue
        _, right = token.split(":", 1)
        cleaned = right.strip().rstrip(",")
        if PROVIDING_PATTERN.match(cleaned):
            facilities.append(cleaned)
    return facilities


def _parse_impact_conditions(tokens: Iterable[str]) -> List[str]:
    conditions: List[str] = []
    for token in tokens:
        if ":" in token:
            _, right = token.split(":", 1)
            cleaned = token.strip().rstrip(",")
            if PROVIDING_PATTERN.match(right.strip().rstrip(",")):
                continue
            conditions.append(cleaned)
    return conditions


def parse_restrictions(raw: str) -> ParsedRestrictions:
    restrictions: List[Restriction] = []
    min_mit: Optional[int] = None
    max_mit: Optional[int] = None
    min_minit: Optional[int] = None
    max_minit: Optional[int] = None
    requesting_facilities: Set[str] = set()
    providing_facilities: Set[str] = set()
    impact_conditions: Set[str] = set()
    timeframes: Set[str] = set()

    for line in raw.splitlines():
        normalized = line.strip()
        if not normalized:
            continue

        tokens = normalized.split()
        timeframe = _match_first(TIMEFRAME_PATTERN, normalized)
        if timeframe:
            timeframes.add(timeframe)

        mit_value = _match_int(MIT_PATTERN, normalized)
        minit_value = _match_int(MINIT_PATTERN, normalized)

        if mit_value is not None:
            min_mit = _update_min(min_mit, mit_value)
            max_mit = _update_max(max_mit, mit_value)
        if minit_value is not None:
            min_minit = _update_min(min_minit, minit_value)
            max_minit = _update_max(max_minit, minit_value)

        requesting = _extract_requesting_facility(normalized)
        providing = _parse_providing_facilities(tokens)
        impacts = _parse_impact_conditions(tokens)

        if requesting:
            requesting_facilities.add(requesting)
        providing_facilities.update(providing)
        impact_conditions.update(impacts)

        restrictions.append(
            Restriction(
                raw=normalized,
                requesting_facility=requesting,
                providing_facilities=providing,
                timeframe=timeframe,
                mit_value=mit_value,
                minit_value=minit_value,
                impact_conditions=impacts,
            )
        )

    return ParsedRestrictions(
        restrictions=restrictions,
        min_mit=min_mit,
        max_mit=max_mit,
        min_minit=min_minit,
        max_minit=max_minit,
        requesting_facilities=requesting_facilities,
        providing_facilities=providing_facilities,
        impact_conditions=impact_conditions,
        timeframes=timeframes,
    )


def _match_first(pattern: re.Pattern[str], text: str) -> Optional[str]:
    match = pattern.search(text)
    return match.group(1) if match else None


def _match_int(pattern: re.Pattern[str], text: str) -> Optional[int]:
    match = pattern.search(text)
    if match:
        try:
            return int(match.group(1))
        except (ValueError, TypeError):
            return None
    return None


def _update_min(current: Optional[int], candidate: int) -> int:
    return candidate if current is None or candidate < current else current


def _update_max(current: Optional[int], candidate: int) -> int:
    return candidate if current is None or candidate > current else current