import cudf
from cuml.neighbors.nearest_neighbors import NearestNeighbors

from ahah.common.logger import logger
from ahah.common.utils import (
    Config,
    clean_hospitals,
    clean_bluespace,
    clean_dentists,
    clean_gpp,
    clean_greenspace_access,
    clean_pharmacies,
    clean_postcodes,
)


def nearest_nodes(df: cudf.DataFrame, nodes: cudf.DataFrame) -> cudf.DataFrame:
    """
    Find nearest road node to point of interest

    Uses `cuml` nearest neighbours for GPU accelerated nearest points.
    It is assumed that all points use a planar coordinate system like BNG.

    Parameters
    ----------
    df : cudf.DataFrame
        POI df containing coordinate information
    nodes : cudf.DataFrame
        Road nodes with coordinate information

    Returns
    -------
    cudf.DataFrame:
        Road nodes that are nearest neighbour to some POI
    """
    df.reset_index(drop=True, inplace=True)
    nodes.reset_index(drop=True, inplace=True)

    points = df[["easting", "northing"]]
    node_points = nodes[["easting", "northing"]]

    nbrs = NearestNeighbors(n_neighbors=1, output_type="cudf").fit(node_points)
    _, indices = nbrs.kneighbors(points)

    nodes = nodes.iloc[indices].reset_index(drop=True)
    return nodes.join(df.drop(["easting", "northing"], axis=1))  # type:ignore


def get_buffers(
    poi: cudf.DataFrame,
    postcodes: cudf.DataFrame,
    k: int,
) -> cudf.DataFrame:
    """
    Estimate buffer sizes required to capture each necessary road node

    Calculates k nearest neighbours for each POI to each road node. Finds
    each node that is considered a neighbour to a poi `k*len(poi)`. Buffers
    are taken as the distance to the further neighbour and all nodes associated with
    each POI are saved.

    Parameters
    ----------
    poi : cudf.DataFrame
        Dataframe of all POIs
    postcodes : cudf.DataFrame
        Dataframe of postcodes
    k : int
        Number of neigbours to use

    Returns
    -------
    cudf.DataFrame:
        POI dataframe including buffer and column with list of nodes
    """
    poi.reset_index(drop=True, inplace=True)
    postcodes.reset_index(drop=True, inplace=True)

    poi_pts = poi[["easting", "northing"]]
    pc_pts = postcodes[["easting", "northing"]]

    nbrs = NearestNeighbors(n_neighbors=k, output_type="cudf").fit(poi_pts)
    distances, indices = nbrs.kneighbors(pc_pts)

    poi = (
        postcodes.join(indices)[["node_id"] + indices.columns.tolist()]
        .set_index("node_id")  # type:ignore
        .stack()
        .rename("node_id")
        .reset_index()
        .rename(columns={0: "pc_ids"})
        .drop(1, axis=1)
        .groupby("node_id")
        .agg(list)
        .join(poi, how="right")
    )
    # retain only unique postcode ids
    poi["pc_ids"] = (
        poi["pc_ids"]
        .to_pandas()
        .apply(lambda row: list(set(row)) if row is not None else row)
    )

    distances = distances.stack().rename("dist").reset_index().drop(1, axis=1)
    indices = indices.stack().rename("ind").reset_index().drop(1, axis=1)

    poi_nodes = (
        poi[["node_id"]].iloc[indices["ind"].values]["node_id"].reset_index(drop=True)
    )
    buffers = cudf.DataFrame({"node_id": poi_nodes, "buffer": distances["dist"].values})
    buffers = buffers.sort_values("buffer", ascending=False).drop_duplicates("node_id")
    buffers["buffer"] = buffers["buffer"].astype("int")

    # this will drops rows that did not appear in the KNN i.e unneeded poi
    return poi.merge(buffers, on="node_id", how="left").dropna()


if __name__ == "__main__":
    logger.info("Starting routing data processing...")
    logger.debug("Reading and cleaning data...")

    nodes = cudf.read_parquet(Config.OSM_GRAPH / "nodes.parquet")
    all_pc: cudf.DataFrame = clean_postcodes(
        path=Config.RAW_DATA / "onspd" / "postcodes.csv", current=False
    )
    cur_pc: cudf.DataFrame = clean_postcodes(
        path=Config.RAW_DATA / "onspd" / "postcodes.csv", current=True
    )
    hospitals: cudf.DataFrame = clean_hospitals(
        england=Config.RAW_DATA / "nhs" / "ets.csv",
        scotland=Config.RAW_DATA / "nhs" / "scotland" / "hospitals.csv",
        postcodes=all_pc,
    )
    gpp: cudf.DataFrame = clean_gpp(
        england=Config.RAW_DATA / "nhs" / "epraccur.csv",
        scotland=Config.RAW_DATA / "nhs" / "scotland" / "gpp.csv",
        postcodes=all_pc,
    )
    dentists: cudf.DataFrame = clean_dentists(
        england=Config.RAW_DATA / "nhs" / "egdpprac.csv",
        scotland=Config.RAW_DATA / "nhs" / "scotland" / "dentists.csv",
        postcodes=all_pc,
    )
    pharmacies: cudf.DataFrame = clean_pharmacies(
        england=Config.RAW_DATA / "nhs" / "edispensary.csv",
        scotland=Config.RAW_DATA / "nhs" / "scotland" / "pharmacies.csv",
        postcodes=all_pc,
    )
    greenspace: cudf.DataFrame = clean_greenspace_access(
        Config.RAW_DATA / "greenspace/access.shp"
    )
    bluespace: cudf.DataFrame = clean_bluespace(Config.RAW_DATA / "bluespace")

    logger.debug("Finding nearest node to postcodes...")
    cur_pc = nearest_nodes(df=cur_pc, nodes=nodes)
    cur_pc.to_parquet(Config.PROCESSED_DATA / "postcodes.parquet", index=False)

    # TODO: uncomment when A&E added
    poi_list = {
        "hospitals": hospitals,
        "gpp": gpp,
        "dentists": dentists,
        "pharmacies": pharmacies,
        "greenspace": greenspace,
        "bluespace": bluespace,
    }
    assert len(poi_list) == len(Config.POI_LIST)

    for poi, df in poi_list.items():
        logger.debug(f"Finding nearest node to {poi}...")
        df = nearest_nodes(df=df, nodes=nodes)
        logger.debug(f"Getting buffers:{poi}...")
        df = get_buffers(poi=df, postcodes=cur_pc, k=10)
        logger.debug(f"Saving {poi}:{Config.PROCESSED_DATA}...")
        df.to_parquet(Config.PROCESSED_DATA / f"{poi}.parquet", index=False)
