metro/excel_graph_loader.py - master - Flash_A/SubjectTemplate

metro/excel_graph_loader.py
259 строк9 KB
FlashAnton
Добавлены материалы от студентов
22 янв 2026, 21:16
22 янв 2026, 21:16511daf0
import os
import re
from functools import lru_cache
from pathlib import Path
from typing import Dict, List, Tuple, Any, Optional

from openpyxl import load_workbook


BASE_HOUR = 5  # 05:00 local start for график
EXCEL_FILENAME = "Сокольническая рабочий осн.xlsx"
STATIONS = [
    "депо Черкизово",
    "Бульвар Рокоссовского",
    "Черкизовская",
    "Преображенская площадь",
    "Сокольники",
    "Красносельская",
    "Комсомольская",
    "Красные ворота",
    "Чистые пруды",
    "Лубянка",
    "Охотный ряд",
    "Библиотека имени Ленина",
    "Кропоткинская",
    "Парк Культуры",
    "Фрунзенская",
    "Спортивная",
    "Воробьёвы горы",
    "Университет",
    "Проспект Вернадского",
    "Юго-Западная",
    "Тропарёво",
    "Румянцево",
    "Саларьево",
    "депо Саларьево",
]

# Mapping station names from Excel to internal names if they differ slighty
STATION_MAPPING = {s: s for s in STATIONS}


class ExcelGraphError(RuntimeError):
    """Raised when excel data cannot be parsed."""



# The list of stations, can be imported by other modules
STATIONS_LIST = STATIONS

def _read_excel_rows(file_source: Any) -> Tuple[List[str], List[List[Any]]]:
    workbook = load_workbook(file_source, data_only=True)
    sheet = workbook.active
    rows = list(sheet.iter_rows(values_only=True))
    if not rows:
        raise ExcelGraphError(f"Загруженный файл не содержит данных")
    headers = [cell if cell is not None else f"col_{idx}" for idx, cell in enumerate(rows[0])]
    return headers, rows[1:]


def _normalize_number(value: Any) -> float:
    if value is None:
        return 0.0
    if isinstance(value, (int, float)):
        return float(value)
    if isinstance(value, str):
        stripped = value.strip().replace(" ", "")
        if not stripped:
            return 0.0
        stripped = stripped.replace(",", ".")
        try:
            return float(stripped)
        except ValueError as exc:
            raise ExcelGraphError(f"Не удалось преобразовать '{value}' в число") from exc
    # Ignore other types (return 0.0) or raise error? For safety, return 0.0
    return 0.0


def _parse_hour_labels(headers: List[str]) -> List[str]:
    hour_labels = []
    for header in headers:
        if isinstance(header, str) and "-" in header:
            start, _, end = header.partition("-")
            if start.isdigit() and end.isdigit():
                hour_labels.append(header)
    if not hour_labels:
        raise ExcelGraphError("В таблице не найдены колонки с часами")
    return hour_labels


def _row_to_dict(headers: List[str], row: List[Any]) -> Dict[str, Any]:
    return {headers[idx]: row[idx] if idx < len(row) else None for idx in range(len(headers))}


def _hour_start_minutes(label: str) -> Tuple[int, int]:
    start_str, _, end_str = label.partition("-")
    start = int(start_str)
    end = int(end_str)
    while start < BASE_HOUR:
        start += 24
    while end <= start:
        end += 24
    start_minutes = (start - BASE_HOUR) * 60
    end_minutes = (end - BASE_HOUR) * 60
    return start_minutes, end_minutes


def _format_time(offset_minutes: float) -> str:
    total_minutes = BASE_HOUR * 60 + offset_minutes
    total_minutes = total_minutes % (24 * 60)
    hours = int(total_minutes // 60)
    minutes = int(round(total_minutes % 60))
    return f"{hours:02d}:{minutes:02d}"


def _generate_runs(
    counts: List[int], hour_labels: List[str], travel_minutes: float, prefix: str
) -> List[Dict[str, Any]]:
    runs: List[Dict[str, Any]] = []
    counter = 1
    for idx, hour_label in enumerate(hour_labels):
        trips = counts[idx]
        if not trips:
            continue
        start_min, end_min = _hour_start_minutes(hour_label)
        interval = 60 / trips
        for trip_idx in range(trips):
            departure_min = start_min + (trip_idx + 0.5) * interval
            arrival_min = departure_min + travel_minutes
            runs.append(
                {
                    "train_id": f"{prefix}-{counter:03d}",
                    "departure_minutes": round(departure_min, 2),
                    "arrival_minutes": round(arrival_min, 2),
                    "departure_time": _format_time(departure_min),
                    "arrival_time": _format_time(arrival_min),
                }
            )
            counter += 1
    return runs


def _parse_route_stations(label: str) -> Tuple[Optional[str], Optional[str]]:
    if not isinstance(label, str) or " - " not in label:
        return None, None
    parts = label.split(" - ")
    if len(parts) != 2:
        return None, None
    
    start, end = parts[0].strip(), parts[1].strip()
    
    # Check if stations exist in our list
    if start not in STATIONS or end not in STATIONS:
        # Try to match loosely if strictly failed? For now strict
        return None, None
        
    return start, end


def _calculate_travel_time(start: str, end: str, row_travel_time: float) -> float:
    if row_travel_time > 0:
        return row_travel_time
    
    # Heuristic calculation based on number of stops
    try:
        start_idx = STATIONS.index(start)
        end_idx = STATIONS.index(end)
        stops = abs(end_idx - start_idx)
        # Assuming ~2.2 minutes per stop on average (including dwell time)
        return stops * 2.2
    except ValueError:
        return 51.0 # Fallback default


def _timeline_meta(hour_labels: List[str]) -> Dict[str, Any]:
    start_min, _ = _hour_start_minutes(hour_labels[0])
    _, final_end = _hour_start_minutes(hour_labels[-1])
    buffer_minutes = 60  # захватываем прибытия после последнего часа
    duration = final_end - start_min + buffer_minutes
    return {
        "base_hour": BASE_HOUR,
        "duration_minutes": duration,
        "start_label": hour_labels[0],
        "end_label": hour_labels[-1],
    }


def load_excel_graph_data(file_source: Any) -> Dict[str, Any]:
    """Загружает данные графика из источника (файл или путь)."""
    headers, raw_rows = _read_excel_rows(file_source)
    source_name = getattr(file_source, 'name', 'uploaded_file')
    hour_labels = _parse_hour_labels(headers)
    first_col_name = headers[0]
    dict_rows = [_row_to_dict(headers, row) for row in raw_rows]

    directions = []
    
    for row in dict_rows:
        label = row.get(first_col_name)
        if not label:
            continue
            
        start_station, end_station = _parse_route_stations(str(label))
        if not start_station or not end_station:
            continue
            
        # Determine direction and indices
        start_idx = STATIONS.index(start_station)
        end_idx = STATIONS.index(end_station)
        
        if start_idx < end_idx:
             # Downwards (e.g. Bukvar Rokossovskogo -> Salaryevo)
             code = "S" # Southbound
             color = "#1f5bd8" # Blue
        else:
             # Upwards (e.g. Salaryevo -> Bulvar Rokossovskogo)
             code = "N" # Northbound
             color = "#0a8e2a" # Green
             
        # Extract counts
        counts = []
        has_trains = False
        for hour_label in hour_labels:
            val = _normalize_number(row.get(hour_label))
            counts.append(int(val))
            if val > 0:
                has_trains = True
                
        if not has_trains:
            continue

        # Extract travel time
        raw_travel_time = _normalize_number(row.get("Тх в пик") or row.get("Тх в непик"))
        travel_minutes = _calculate_travel_time(start_station, end_station, raw_travel_time)
        
        runs = _generate_runs(counts, hour_labels, travel_minutes, f"{code}_{start_idx}_{end_idx}")
        
        directions.append({
            "name": label,
            "start_station": start_station,
            "end_station": end_station,
            "start_idx": start_idx,
            "end_idx": end_idx,
            "color": color,
            "code": code,
            "travel_minutes": travel_minutes,
            "hour_counts": counts,
            "total_departures": sum(counts),
            "runs": runs,
        })

    payload = {
        "source_file": source_name,
        "stations": STATIONS,
        "hours": hour_labels,
        "timeline": _timeline_meta(hour_labels),
        "directions": directions,
    }
    return payload