from typing import List

import numpy as np
import pandas as pd
from TrainInfoProcess import TrainNumberInfo, originStations


def TrainNumberInfosExtraction(database: pd.DataFrame) -> List[TrainNumberInfo]:
    """提取列车信息

    Args:
        database (pd.DataFrame): 经初始化后的原始数据集

    Returns:
        List[TrainNumberInfo]: 列车信息
    """
    trainNumberInfos = []
    for trainNumber in database["id"].unique():
        data = database.loc[database["id"] == trainNumber]
        originStation = data["origin"].iloc[0]
        destination = data["destination"].iloc[0]
        trainNumberInfo = TrainNumberInfo(trainNumber, originStation, destination)
        for station in originStations:
            arrivalDf = data.loc[data["to"] == station]
            departureDf = data.loc[data["from"] == station]
            if arrivalDf.empty:
                arrivalTime = np.nan
            else:
                arrivalTime = (
                    arrivalDf["time4to"].iloc[0].hour * 60
                    + arrivalDf["time4to"].iloc[0].minute
                )
            if departureDf.empty:
                departureTime = np.nan
            else:
                departureTime = (
                    departureDf["time4from"].iloc[0].hour * 60
                    + departureDf["time4from"].iloc[0].minute
                )
            flag = arrivalDf.empty and departureDf.empty
            if flag:
                continue
            trainNumberInfo.addStation(station, arrivalTime, departureTime)
        trainNumberInfo.sortStations()
        trainNumberInfos.append(trainNumberInfo)

    return trainNumberInfos

from typing import List

import numpy as np
import pandas as pd
from TrainInfoProcess import TrainNumberInfo, originStations


def TrainNumberInfosExtraction(database: pd.DataFrame) -> List[TrainNumberInfo]:
    """提取列车信息

    Args:
        database (pd.DataFrame): 经初始化后的原始数据集

    Returns:
        List[TrainNumberInfo]: 列车信息
    """
    trainNumberInfos = []
    for trainNumber in database["id"].unique():
        data = database.loc[database["id"] == trainNumber]
        originStation = data["origin"].iloc[0]
        destination = data["destination"].iloc[0]
        trainNumberInfo = TrainNumberInfo(trainNumber, originStation, destination)
        for station in originStations:
            arrivalDf = data.loc[data["to"] == station]
            departureDf = data.loc[data["from"] == station]
            if arrivalDf.empty:
                arrivalTime = np.nan
            else:
                arrivalTime = (
                    arrivalDf["time4to"].iloc[0].hour * 60
                    + arrivalDf["time4to"].iloc[0].minute
                )
            if departureDf.empty:
                departureTime = np.nan
            else:
                departureTime = (
                    departureDf["time4from"].iloc[0].hour * 60
                    + departureDf["time4from"].iloc[0].minute
                )
            flag = arrivalDf.empty and departureDf.empty
            if flag:
                continue
            trainNumberInfo.addStation(station, arrivalTime, departureTime)
        trainNumberInfo.sortStations()
        trainNumberInfos.append(trainNumberInfo)

    return trainNumberInfos

提取列车信息      Args:         database (pd.DataFrame): 经初始化后的原始数据集      Returns:         List[TrainNumberInfo]: 列车信息