from __future__ import annotations

import logging
import math

from typing import TYPE_CHECKING

import networkx as nx

from pyproj import Transformer

if TYPE_CHECKING:
    from networkx import Graph

    from edisgo import EDisGo

logger = logging.getLogger(__name__)

# Transform coordinates to equidistant and back
coor_transform = Transformer.from_crs("EPSG:4326", "EPSG:3035", always_xy=True)
coor_transform_back = Transformer.from_crs("EPSG:3035", "EPSG:4326", always_xy=True)


# Pseudo coordinates
def _make_coordinates(graph_root: Graph, branch_detour_factor: float) -> Graph:
    """
    Generates pseudo coordinates for a graph with equidistant coordinates.

    Parameters
    ----------
    graph_root : :networkx:`networkx.Graph<>`
        Graph object to generate pseudo coordinates for (with equidistant coordinates).

    branch_detour_factor : float
        Defines the quotient of the line length and the distance of the buses.

    Returns
    -------
    :networkx:`networkx.Graph<>`
        Graph with equidistant pseudo coordinates for all nodes.

    """

    # Make coordinates for the neighbours of the transformer node
    # Nodes are distributed around the source node with the same angle
    def coordinate_source(pos_start, length, node_numerator, node_total_numerator):
        length = length / branch_detour_factor
        angle = node_numerator * 360 / node_total_numerator
        x0, y0 = pos_start
        x1 = x0 + 1000 * length * math.cos(math.radians(angle))
        y1 = y0 + 1000 * length * math.sin(math.radians(angle))
        pos_end = (x1, y1)
        origin_angle = math.degrees(math.atan2(y1 - y0, x1 - x0))
        return pos_end, origin_angle

    # Make coordinates for nodes which are not neighbors of the transformer,
    # not in the longest path, not neighbors of the longest path.
    # Nodes are placed in the half plane generated by the straight longest path,
    # the angle between the neighbouring nodes are same, basically a tree.
    def coordinate_branch(
        pos_start, angle_offset, length, node_numerator, node_total_numerator
    ):
        length = length / branch_detour_factor
        angle = node_numerator * 180 / (node_total_numerator + 1) + angle_offset - 90
        x0, y0 = pos_start
        x1 = x0 + 1000 * length * math.cos(math.radians(angle))
        y1 = y0 + 1000 * length * math.sin(math.radians(angle))
        origin_angle = math.degrees(math.atan2(y1 - y0, x1 - x0))
        pos_end = (x1, y1)
        return pos_end, origin_angle

    # Make coordinates for the nodes of the longest path
    # Nodes are distributed in a straight line
    def coordinate_longest_path(pos_start, angle_offset, length):
        length = length / branch_detour_factor
        angle = angle_offset
        x0, y0 = pos_start
        x1 = x0 + 1000 * length * math.cos(math.radians(angle))
        y1 = y0 + 1000 * length * math.sin(math.radians(angle))
        origin_angle = math.degrees(math.atan2(y1 - y0, x1 - x0))
        pos_end = (x1, y1)
        return pos_end, origin_angle

    # Make coordinates for neighbours of nodes of the longest path
    # Nodes are placed in an angle of 90 degrees to the longest path
    # with alternating direction
    def coordinate_longest_path_neighbor(pos_start, angle_offset, length, direction):
        length = length / branch_detour_factor
        if direction:
            angle_random_offset = 90
        else:
            angle_random_offset = -90
        angle = angle_offset + angle_random_offset
        x0, y0 = pos_start
        x1 = x0 + 1000 * length * math.cos(math.radians(angle))
        y1 = y0 + 1000 * length * math.sin(math.radians(angle))
        origin_angle = math.degrees(math.atan2(y1 - y0, x1 - x0))
        pos_end = (x1, y1)

        return pos_end, origin_angle

    # Find transformer node and copy graph
    start_node = list(nx.nodes(graph_root))[0]
    graph_root.nodes[start_node]["pos"] = (0, 0)
    graph_copy = graph_root.copy()

    long_paths = []
    next_nodes = []

    # Find longest paths
    for i in range(0, len(list(nx.neighbors(graph_root, start_node)))):
        path_length_to_transformer = []
        for node in graph_copy.nodes():
            try:
                paths = list(nx.shortest_simple_paths(graph_copy, start_node, node))
            except nx.NetworkXNoPath:
                paths = [[]]
            path_length_to_transformer.append(len(paths[0]))
        index = path_length_to_transformer.index(max(path_length_to_transformer))
        path_to_max_distance_node = list(
            nx.shortest_simple_paths(
                graph_copy, start_node, list(nx.nodes(graph_copy))[index]
            )
        )[0]
        path_to_max_distance_node.remove(start_node)
        graph_copy.remove_nodes_from(path_to_max_distance_node)
        for node in path_to_max_distance_node:
            long_paths.append(node)

    path_to_max_distance_node = long_paths
    n = 0

    # make the coordinates
    for node in list(nx.neighbors(graph_root, start_node)):
        n = n + 1
        pos, origin_angle = coordinate_source(
            graph_root.nodes[start_node]["pos"],
            graph_root.edges[start_node, node]["length"],
            n,
            len(list(nx.neighbors(graph_root, start_node))),
        )
        graph_root.nodes[node]["pos"] = pos
        graph_root.nodes[node]["origin_angle"] = origin_angle
        next_nodes.append(node)

    graph_copy = graph_root.copy()
    graph_copy.remove_node(start_node)
    while graph_copy.number_of_nodes() > 0:
        next_node = next_nodes[0]
        n = 0
        for node in list(nx.neighbors(graph_copy, next_node)):
            n = n + 1
            if node in path_to_max_distance_node:
                pos, origin_angle = coordinate_longest_path(
                    graph_root.nodes[next_node]["pos"],
                    graph_root.nodes[next_node]["origin_angle"],
                    graph_root.edges[next_node, node]["length"],
                )
            elif next_node in path_to_max_distance_node:
                direction = math.fmod(
                    len(
                        list(
                            nx.shortest_simple_paths(graph_root, start_node, next_node)
                        )[0]
                    ),
                    2,
                )
                pos, origin_angle = coordinate_longest_path_neighbor(
                    graph_root.nodes[next_node]["pos"],
                    graph_root.nodes[next_node]["origin_angle"],
                    graph_root.edges[next_node, node]["length"],
                    direction,
                )
            else:
                pos, origin_angle = coordinate_branch(
                    graph_root.nodes[next_node]["pos"],
                    graph_root.nodes[next_node]["origin_angle"],
                    graph_root.edges[next_node, node]["length"],
                    n,
                    len(list(nx.neighbors(graph_copy, next_node))),
                )

            graph_root.nodes[node]["pos"] = pos
            graph_root.nodes[node]["origin_angle"] = origin_angle
            next_nodes.append(node)

        graph_copy.remove_node(next_node)
        next_nodes.remove(next_node)

    return graph_root


def make_pseudo_coordinates_graph(G: Graph, branch_detour_factor: float) -> Graph:
    """
    Generates pseudo coordinates for one graph.

    Parameters
    ----------
    G : :networkx:`networkx.Graph<>`
        Graph object to generate pseudo coordinates for.

    branch_detour_factor : float
        Defines the quotient of the line length and the distance of the buses.

    Returns
    -------
    :networkx:`networkx.Graph<>`
        Graph with pseudo coordinates for all nodes.

    """
    x0, y0 = G.nodes[list(nx.nodes(G))[0]]["pos"]
    G = _make_coordinates(G, branch_detour_factor)
    x0, y0 = coor_transform.transform(x0, y0)
    for node in G.nodes():
        x, y = G.nodes[node]["pos"]
        G.nodes[node]["pos"] = coor_transform_back.transform(x + x0, y + y0)
    return G


def make_pseudo_coordinates(edisgo_obj: EDisGo, mv_coordinates: bool = False):
    """
    Generates pseudo coordinates for all LV grids and optionally MV grid.

    Bus coordinates are changed in the Topology object directly. If you want to keep
    information on the original coordinates, hand a copy of the EDisGo object to this
    function.

    Parameters
    ----------
    edisgo_obj : :class:`~.EDisGo`
        eDisGo object to create pseudo coordinates for.
    mv_coordinates : bool, optional
        If False, pseudo coordinates are only generated for LV buses. If True, pseudo
        coordinates are as well generated for MV buses.
        Default: False.

    """
    grids = list(edisgo_obj.topology.mv_grid.lv_grids)
    if mv_coordinates:
        grids = [edisgo_obj.topology.mv_grid] + grids

    for grid in grids:
        logger.debug("Make pseudo coordinates for: {}".format(grid))
        G = grid.graph
        G = make_pseudo_coordinates_graph(
            G, edisgo_obj.config["grid_connection"]["branch_detour_factor"]
        )
        for node in G.nodes():
            x, y = G.nodes[node]["pos"]
            edisgo_obj.topology.buses_df.loc[node, "x"] = x
            edisgo_obj.topology.buses_df.loc[node, "y"] = y
