import os

import numpy as np
import pandas as pd
import torch


def get_adj(args):
    dataset_path = "./data"
    match args["num_nodes"]:
        case 358:
            dataset_name = "PEMS03"
            adj_path = os.path.join(dataset_path, dataset_name, "PEMS03.csv")
            id = os.path.join(dataset_path, dataset_name, "PEMS03.txt")
            A = get_adjacency_matrix(
                adj_path, args["num_nodes"], args["construct_type"], id_filename=id
            )
        case 307:
            dataset_name = "PEMS04"
            adj_path = os.path.join(dataset_path, dataset_name, "PEMS04.csv")
            A = get_adjacency_matrix(
                adj_path, args["num_nodes"], args["construct_type"]
            )
        case 883:
            dataset_name = "PEMS07"
            adj_path = os.path.join(dataset_path, dataset_name, "PEMS07.csv")
            A = get_adjacency_matrix(
                adj_path, args["num_nodes"], args["construct_type"]
            )
        case 170:
            dataset_name = "PEMS08"
            adj_path = os.path.join(dataset_path, dataset_name, "PEMS08.csv")
            A = get_adjacency_matrix(
                adj_path, args["num_nodes"], args["construct_type"]
            )

    local_adj = construct_adj(A, args["strides"])
    local_adj = torch.FloatTensor(local_adj)

    return local_adj


def get_adjacency_matrix(
    distance_df_filename, num_of_vertices, type_="connectivity", id_filename=None
):
    """
    :param distance_df_filename: str, csv边信息文件路径
    :param num_of_vertices:int, 节点数量
    :param type_:str, {connectivity, distance}
    :param id_filename:str 节点信息文件， 有的话需要构建字典
    """
    A = np.zeros((int(num_of_vertices), int(num_of_vertices)), dtype=np.float32)

    if id_filename:
        with open(id_filename, "r") as f:
            id_dict = {
                int(i): idx for idx, i in enumerate(f.read().strip().split("\n"))
            }  # 建立映射列表
        df = pd.read_csv(distance_df_filename)
        for row in df.values:
            if len(row) != 3:
                continue
            i, j = int(row[0]), int(row[1])
            A[id_dict[i], id_dict[j]] = 1
            A[id_dict[j], id_dict[i]] = 1

        return A

    df = pd.read_csv(distance_df_filename)
    for row in df.values:
        if len(row) != 3:
            continue
        i, j, distance = int(row[0]), int(row[1]), float(row[2])
        if type_ == "connectivity":
            A[i, j] = 1
            A[j, i] = 1
        elif type == "distance":
            A[i, j] = 1 / distance
            A[j, i] = 1 / distance
        else:
            raise ValueError("type_ error, must be connectivity or distance!")

    return A


def construct_adj(A, steps):
    """
    构建local 时空图
    :param A: np.ndarray, adjacency matrix, shape is (N, N)
    :param steps: 选择几个时间步来构建图
    :return: new adjacency matrix: csr_matrix, shape is (N * steps, N * steps)
    """
    N = len(A)  # 获得行数
    adj = np.zeros((N * steps, N * steps))

    for i in range(steps):
        """对角线代表各个时间步自己的空间图，也就是A"""
        adj[i * N : (i + 1) * N, i * N : (i + 1) * N] = A

    for i in range(N):
        for k in range(steps - 1):
            """每个节点只会连接相邻时间步的自己"""
            adj[k * N + i, (k + 1) * N + i] = 1
            adj[(k + 1) * N + i, k * N + i] = 1

    for i in range(len(adj)):
        """加入自回"""
        adj[i, i] = 1

    return adj