import os
import pandas as pd
from pandas import DataFrame


def time2float(time_str):
    h, m = time_str.split(":")
    return int(h) + int(m) / 60


class IntercityTransport:
    def __init__(self, path: str = "../../database/intercity_transport/"):
        curdir = os.path.dirname(os.path.realpath(__file__))
        self.base_path = os.path.join(curdir, path)
        self.airplane_path = self.base_path + "airplane.jsonl"
        # self.airplane_path = self.base_path + "airplane_real_2025_08_13.jsonl"
        self.airplane_df = pd.read_json(
            self.airplane_path, lines=True, keep_default_dates=False
        )
        city_list = [
            "上海",
            "北京",
            "深圳",
            "广州",
            "重庆",
            "苏州",
            "成都",
            "杭州",
            "武汉",
            "南京",
        ]
        self.train_df_dict = {}

        for start_city in city_list:
            for end_city in city_list:
                if start_city == end_city:
                    continue
                train_path = (
                    self.base_path
                    + "train/"
                    + "from_{}_to_{}.json".format(start_city, end_city)
                )
                train_df = pd.read_json(train_path)
                self.train_df_dict[(start_city, end_city)] = train_df

    def select(
        self, start_city, end_city, intercity_type, earliest_leave_time="00:00"
    ) -> DataFrame:
        if intercity_type not in ["train", "airplane"]:
            return "only support intercity_type in ['train','airplane']"
        res = self._select(start_city, end_city, intercity_type)
        bool_list = [False] * len(res)
        for i in range(len(res)):
            if time2float(res.loc[i, "BeginTime"]) >= time2float(earliest_leave_time):
                bool_list[i] = True
        return res[bool_list]

    def _select(self, start_city, end_city, intercity_type) -> DataFrame:
        # intercity_type=='train' | 'airplane'
        if intercity_type == "airplane":

            if len(self.airplane_df) == 0:
                return None

            filtered_flights = self.airplane_df[
                (self.airplane_df["From"].str.contains(start_city))
                & (self.airplane_df["To"].str.contains(end_city))
            ]
            sorted_flights = filtered_flights.sort_values(by="BeginTime").reset_index(
                drop=True
            )
            return sorted_flights
        if intercity_type == "train":

            if len(self.train_df_dict[(start_city, end_city)]) == 0:
                return None

            filtered_trains = self.train_df_dict[(start_city, end_city)]
            sorted_trains = filtered_trains.sort_values(by="BeginTime").reset_index(
                drop=True
            )
            return sorted_trains


if __name__ == "__main__":
    a = IntercityTransport()
    city_list = [
        "上海",
        "北京",
        "深圳",
        "广州",
        "重庆",
        "苏州",
        "成都",
        "杭州",
        "武汉",
        "南京",
    ]
    city_en_list = [
        "shanghai",
        "beijing",
        "shenzhen",
        "guangzhou",
        "chongqing",
        "suzhou",
        "chengdu",
        "hangzhou",
        "wuhan",
        "nanjing",
    ]
    str_list = []
    for i in range(len(city_list)):
        for j in range(i + 1, len(city_list)):
            tmp_len = 0

            tmp = a.select(city_list[i], city_list[j], "train")
            if not isinstance(tmp, DataFrame):
                tmp_len = 0
            else:
                tmp_len = len(tmp)
            if tmp_len > 0:
                str_list.append(
                    "('{}','{}','{}',{})".format(
                        city_en_list[i], city_en_list[j], "train", tmp_len
                    )
                )

            tmp = a.select(city_list[j], city_list[i], "train")
            if not isinstance(tmp, DataFrame):
                tmp_len = 0
            else:
                tmp_len = len(tmp)
            if tmp_len > 0:
                str_list.append(
                    "('{}','{}','{}',{})".format(
                        city_en_list[j], city_en_list[i], "train", tmp_len
                    )
                )

            tmp = a.select(city_list[i], city_list[j], "flight")
            if not isinstance(tmp, DataFrame):
                tmp_len = 0
            else:
                tmp_len = len(tmp)
            if tmp_len > 0:
                str_list.append(
                    "('{}','{}','{}',{})".format(
                        city_en_list[i], city_en_list[j], "airplane", tmp_len
                    )
                )

            tmp = a.select(city_list[j], city_list[i], "flight")
            if not isinstance(tmp, DataFrame):
                tmp_len = 0
            else:
                tmp_len = len(tmp)
            if tmp_len > 0:
                str_list.append(
                    "('{}','{}','{}',{})".format(
                        city_en_list[j], city_en_list[i], "airplane", tmp_len
                    )
                )

    print(",\n".join(str_list))
