from __future__ import annotations

import copy
import logging
import os
import pickle
from io import TextIOWrapper
from itertools import chain
from typing import Callable

import networkx as nx
# import rustworkx as rx
import numpy as np
from shapely.geometry import Point
from shapely.ops import unary_union

from canopyhydro.configuration import in_flow_grade_lim, input_dir, output_dir
from canopyhydro.Cylinder import Cylinder, create_cyl
from canopyhydro.DataClasses import Flow
from canopyhydro.geometry import (concave_hull, draw_cylinders_3D, draw_cyls,
                                  furthest_point, get_projected_overlap)
from canopyhydro.utils import (create_dir_and_file, intermitent_log,
                               lam_filter, save_file)

log = logging.getLogger("model")

NAME = "CylinderCollection"


def pickle_collection(collection: CylinderCollection, file_name_ext: str = ""):
    """
    Pickles a collection to a file so that it may be
        read into a CylinderCollection object at a later time.

    Args:
        collection:
            The collection to pickle.
        file_name_ext:
            the naming convention to used for the saved file name
            i.e. For a collection w/ file name 'myQSM.csv' and file_name_ext
                'test_run' the file created will be 'myQSM_test_run.pickle'
    Example:
        For a collection w/ file name 'myQSM.csv' and file_name_ext
            'test_run' the file created will be 'myQSM_test_run.pickle'

    Returns:
        file_path (str): The path and file name for the the pickled file.

    Raises:
        FileNotFoundError: If the file does not exist.
        pickle.UnpicklingError: If there is an error while unpickling the file.
    """
    qsm_file = f'{collection.file_name.replace(".csv","")}'
    if ".pickle" not in file_name_ext:
        file_name_ext = f"{file_name_ext}.pickle"
    file_to_create = f"{qsm_file}_{file_name_ext}"
    file_path = "".join([output_dir, "pickle/", file_to_create])
    directory = os.path.dirname(file_path)
    create_dir_and_file(directory)
    pickle_file = open(file_path, "ab")
    pickle.dump(collection, pickle_file)
    pickle_file.close()
    return file_path


def unpickle_collection(file_name: str) -> CylinderCollection:
    """
    Unpickles a collection from a file.
    File assumed to be generated by 'pickle_collection' method

    Args:
        file_name (str):
            The name of the file to unpickle. Optionally, the full path to the file.

    Returns:
        db: The unpickled collection.

    Raises:
        FileNotFoundError: If the file does not exist.
        pickle.UnpicklingError: If there is an error while unpickling the file.
    """
    # In the event that the passed file name is
    if "pickle/" in file_name:
        file_path = file_name
    else:
        file_path = "".join([output_dir, "pickle/", file_name])
    dbfile = open(file_path, "rb")
    db = pickle.load(dbfile)
    dbfile.close()
    return db


class CylinderCollection:
    """
    CylinderCollections represent QSM's or sections thereof, but with additional structure.
    A CylinderCollection is initialized from a .csv file and can:
        - Project cylinders onto the XY, XZ, and YZ planes
        - Determine which precipitation partition water intercepted by
            each of its cylinders will flow into.
        - Find common tree characteristics: DBH, trunk lean, etc.
        - Calculate occlusion for given cross sections of the canopy
        - Draw detailed figures representing the canopy watershed, with:
            - Cylinders in 3D or projected on to the XY, XZ, or YZ planes
            - Drip points (areas of dense throughfall)
            - Alpha shapes (concave hulls) around arbitrary parts of the canopy

    Attributes:
       file (str): The path to the .csv file used to initialize the CylinderCollection. Defaults to "".
       surface_area (float):
          The total surface area of all cylinders in the collection. Defaults to np.nan.
       file_name (string):
          The name of the .csv file used to initialize the CylinderCollection. Defaults to "".
       volume (float):
          The total volume of all cylinders in the collection. Defaults to np.nan.
       avg_sa_to_vol (float):
          The average ratio of surface area to volume for the cylinders in the collection. Defaults to np.nan.
       max_branch_order (float):
          The maximum branch order in the collection. Defaults to np.nan.
       max_reverse_branch_order (float):
          The maximum reverse branch order in the collection. Defaults to np.nan.
       extent (dict):
          The minimum and maximum coordinates of the cylinders in the collection.
          Defaults to {
               "min": [np.nan, np.nan, np.nan],
               "max": [np.nan, np.nan, np.nan],
           }.
       treeQualities (dict):
          Various characteristics of the tree represented by the collection. Defaults to {
           "total_psa": None,
           "tot_hull_area": None,
           "stem_flow_hull_area": None,
           "stem_psa": None,
           "flowStats": None,
           "dbh": None,
           "tot_surface_area": None,
           "stem_surface_area": None,
       }.
       projections (dict):
          Whether the cylinders have been projected onto the XY, XZ, and YZ planes.
          Defaults to {
           "XY": False,
           "XZ": False,
           "YZ": False
           }.
       hulls (dict):
          Whether hulls have been calculated for the projections onto the XY, XZ, and YZ planes.
          Defaults to {
           "XY":False,
           "XZ": False,
           "YZ": False
          }.
       stem_hull (shapely.Polygon):
          The hull of the stem of the tree. Defaults to None.
       stem_paths (list):
          The paths along the stem of the tree. Defaults to [[]].
       trunk (list):
          The cylinders that make up the trunk of the tree. Defaults to [].
       graph (nx.Graph):
          The graph representation of the tree. Defaults to None.
       digraph (nx.Digraph):
          The directed graph representation of the tree. Defaults to None.
       end_nodes (list):
          The nodes at the ends of the branches of the tree. Defaults to [].
       flows (list[Flow]):
          The flows of water intercepted by the cylinders. Defaults to None.
       divide_nodes (list[int]):
          The nodes at which the flows divide. Defaults to None.
       drip_nodes (list[int]):
          The nodes at which water drips from the tree. Defaults to None.
       cyl_to_drip (dict[int,int]):
          A mapping from cylinders to drip points. Defaults to None.
       drip_points (dict[str,list]):
          The coordinates and flow ids of the drip points.
          Defaults to {
           "x": np.nan,
           "y": np.nan,
           "z": np.nan,
           "flow_id": np.nan
           }.
       trunk_nodes (list):
          The nodes that make up the trunk of the tree. Defaults to [].
       drip_point_loc (list):
          The locations of the drip points. Defaults to None.
       stem_flow_component (nx.Graph):
          The component of the graph representing stem flow. Defaults to None.
       drip_graph (nx.Graph):
          The graph representing the flow of drips. Defaults to None.
       drip_flow_components (list[nx.Graph]):
          The components of the drip graph. Defaults to None.
       flow_chars (dict):
          Characteristics of the flows: num_cylinders, projected_area,
           surface_area, angle_sum, volume, sa_to_vol, drip_node_id, drip_node_loc.
          Defaults to {}.
       trunk_lean (float):
          The angle of lean of the trunk. Defaults to None.
    """

    cylinders = dict

    # initialize our object level variables for cylinder objects
    def __init__(self) -> None:
        self.file = ""
        # Aggregate values from file
        self.surface_area = np.nan
        self.file_name = ""
        self.pSV = []  # Whole tree polygon
        self.volume = np.nan
        self.avg_sa_to_vol = np.nan
        self.max_branch_order = np.nan
        self.max_reverse_branch_order = np.nan
        self.extent = {
            "min": [np.nan, np.nan, np.nan],
            "max": [np.nan, np.nan, np.nan],
        }
        # to populate with x,y,z mins and maxs
        self.treeQualities = {
            "total_psa": None,
            "tot_hull_area": None,
            "stem_flow_hull_area": None,
            "stem_psa": None,
            "flowStats": None,
            "dbh": None,
            "tot_surface_area": None,
            "stem_surface_area": None,
        }

        # Projection Attrs
        self.projections = {"XY": False, "XZ": False, "YZ": False}
        self.hulls = {"XY": False, "XZ": False, "YZ": False}
        self.stem_hull = None

        # Special case tree attributes
        self.stem_paths = [[]]
        self.trunk = []

        # Graph and Attributes
        self.graph = None
        self.digraph = None
        self.end_nodes = []
        self.flows = None
        self.divide_nodes = None
        self.drip_nodes = None
        self.cyl_to_drip = None
        self.drip_points = {"x": np.nan, "y": np.nan, "z": np.nan, "flow_id": np.nan}
        self.trunk_nodes = []
        self.drip_point_loc = None
        self.stem_flow_component = None
        self.drip_graph = None
        self.drip_flow_components = None
        # Calculations using graph results
        self.flow_chars = {}
        self.trunk_lean = None

    def __eq__(self, other):
        """
        Compares two CylinderCollections for equality.

        Args:
            other: The other CylinderCollection to compare.

        Returns:
            bool: True if the cylinders in the two collections are the same, False otherwise.

        Raises:
            AttributeError: If one or both CylinderCollections contain no cylinders.
            TypeError: If the other object is not a CylinderCollection.

        """
        if len(self.cylinders) == 0 or len(other.cylinders) == 0:
            raise AttributeError(
                "One or both Cylinder Collections contain no Cylinders (Did you forget to initialize?)"
            )
        if type(self) == type(other):
            raise TypeError(
                "CylinderCollections may only be compared to other cylinder collections"
            )
        # order matters here
        return np.all([cyl == other[idx] for idx, cyl in enumerate(self.cylinders)])

    def from_csv(self, file: TextIOWrapper | str) -> None:
        """
        Initializes a new Cylinder Collection based on the QSM data.
         the csv file is read in with the columns mapped to the configured column locations.

        Args:
            file: The file object or file name to read the data from.

        """
        if isinstance(file, str):
            self.file_name = file
            file = open(input_dir + file)
        else:
            self.file = file
            self.file_name = file.name
        log.info(f"Processing {str(file)}")
        self.arr = np.genfromtxt(file, delimiter=",", skip_header=True)[0:, :-1]
        cylinders = [create_cyl(row) for row in self.arr]
        self.cylinders = cylinders

        min_x = np.min([cyl.x[0] for cyl in cylinders])
        min_y = np.min([cyl.y[0] for cyl in cylinders])
        min_z = np.min([cyl.z[0] for cyl in cylinders])
        max_x = np.max([cyl.x[1] for cyl in cylinders])
        max_y = np.max([cyl.y[1] for cyl in cylinders])
        max_z = np.max([cyl.z[1] for cyl in cylinders])
        # Aggregate values from file
        self.no_cylinders = len(cylinders)
        self.surface_area = np.sum([cyl.surface_area for cyl in cylinders])
        self.volume = np.sum([cyl.volume for cyl in cylinders])
        self.max_branch_order = np.max([cyl.branch_order for cyl in cylinders])
        self.max_reverse_branch_order = np.max(
            [cyl.reverse_branch_order for cyl in cylinders]
        )
        self.avg_sa_to_vol = (
            np.sum([cyl.sa_to_vol for cyl in cylinders]) / self.no_cylinders
        )
        self.extent = {
            "min": [min_x, min_y, min_z],
            "max": [max_x, max_y, max_z],
        }

        self.theta = np.nan
        log.info(f"{file.name} initialized with {self.no_cylinders} cylinders")

    def project_cylinders(self, plane: str = "XY") -> None:
        """
        Projects cylinders onto the specified plane.

        Args:
            plane (str): The plane to project the cylinders onto. Valid values are "XY", "XZ", and "YZ".

        """
        if plane not in ("XY", "XZ", "YZ"):
            log.info(f"{plane}: invalid value for plane")

        else:
            polys = []
            log.info(f"Projection into {plane} axis begun for file {self.file_name}")
            for idx, cyl in enumerate(self.cylinders):
                poly = cyl.get_projection(plane)
                polys.append(poly)
                # print a progress update once every 10 thousand or so cylinders
                intermitent_log(idx, self.no_cylinders, "Cylinder projection: ")
            # Used by other functions to know what projections have been run
            self.projections[plane] = True
            self.pSV = polys
            log.info(f"Projection into {plane} axis complete for file {self.file_name}")

    def get_collection_data(self) -> list[str]:
        """
        Returns a list of string representations of the cylinders in the collection.

        Returns:
            list[str]: A list of string representations of the cylinders.
        """
        cyl_desc = [cyl.__repr__() for cyl in self.cylinders]
        return cyl_desc

    def draw(
        self,
        plane: str = "XY",  # XY, XZ, YZ, or 3D
        highlight_lambda: Callable = lambda: True,
        filter_lambda: Callable = lambda: True,
        include_drips: bool = False,
        include_alpha_shape: bool = False,
        overlay_transparency: float = 0.5,
        file_name_ext: str = "",
        **kwargs,
    ):
        """
        Draws cylinders meeting given characteristics onto the specified plane.

        Args:
            plane:
                The plane to draw the cylinders on. Valid values are "XY", "XZ", "YZ", and "3D".
            highlight_lambda:
                A lambda function that determines which cylinders to highlight.
                Takes no inputs but may reference any attribute of a cylinder object.
            filter_lambda:
                A lambda function that determines which cylinders to include in the drawing.
                Takes no inputs but may reference any attribute of a cylinder object.
            include_drips:
                If True, includes the most recently set drip points in the drawing.
                Default is False.
            include_alpha_shape:
                If True, includes the most recently created alpha shape in the drawing.
                Default is False.
            overlay_transparency:
                Determines the transparency of the drip points and alpha shape drawn.
                Default is 0.5,
            file_name_ext:
                A suffix and/or the file extension to use for the output file. Default is an empty string.
            **kwargs:
            Additional keyword arguments to pass to the drawing function.
        """
        if plane not in ("XY", "XZ", "YZ", "3D"):
            log.info(f"{plane}: invalid value for plane")
        file_name = f"{self.file_name}_{plane}_{file_name_ext}"
        cylinders, _ = lam_filter(self.cylinders, filter_lambda)
        filtered_cyls, matches = lam_filter(
            cylinders, highlight_lambda, return_all=True
        )
        log.info(f"{len(filtered_cyls)} cylinders matched criteria")
        if plane == "3D":
            radii, vectors = [], []
            for cyl in filtered_cyls:
                start_end = np.array(
                    [
                        np.array([cyl.x[0], cyl.y[0], cyl.z[0]]),
                        np.array([cyl.x[1], cyl.y[1], cyl.z[1]]),
                    ]
                )
                radius = cyl.radius
                radii.append(radius)
                vectors.append(start_end)
            fig = draw_cylinders_3D(radii, vectors, file_name_ext=file_name, **kwargs)
        else:
            if not self.projections[plane]:
                self.project_cylinders(plane)
            to_draw = [cyl.projected_data[plane]["polygon"] for cyl in filtered_cyls]

            overlay = []
            if include_alpha_shape:
                if not self.hulls[plane]:
                    log.warning(
                        f"No {plane} plane alpha shape found. Running watershed_boundary"
                    )
                    self.watershed_boundary(plane=plane)
                overlay.append(self.hulls[plane])

            if include_drips:
                if not self.drip_point_loc:
                    log.warning(
                        "No drip point locations found, running set_drip_points"
                    )
                    self.set_drip_points()
                dpts = self.drip_point_loc
                x_pts = list(pt[0] for pt in dpts)
                y_pts = list(pt[1] for pt in dpts)
                z_pts = list(pt[2] for pt in dpts)
                if plane == "XY":
                    overlay.append([x_pts, y_pts])
                elif plane == "XZ":
                    overlay.append([x_pts, z_pts])
                elif plane == "YZ":
                    overlay.append([y_pts, z_pts])

            fig = draw_cyls(
                collection=to_draw,
                colors=matches,
                file_name_ext=file_name,
                overlay=overlay,
                overlay_transparency=overlay_transparency,
                **kwargs,
            )
        return fig

    def get_dbh(self) -> float:
        """
        Calculate and return the diameter at breast height (DBH) of the tree.

        Returns:
            float: The diameter at breast height (DBH) of the tree.
        """
        if self.treeQualities["dbh"]:
            return self.treeQualities["dbh"]
        else:
            desired_height = 1.3
            approx_height_range_start = desired_height - 0.1
            start_z = self.extent["min"][2]

            diff_from_desired_height = [
                (abs(cyl.z[0] - desired_height), cyl.cyl_id)
                for cyl in self.cylinders
                if cyl.z[0] >= start_z + approx_height_range_start
                and cyl.branch_order == 0
            ]
            closest_to_target_id = np.argmin(np.array(diff_from_desired_height[:][0]))
            rbh = self.cylinders[
                int(diff_from_desired_height[closest_to_target_id][1])
            ].radius
            self.treeQualities["dbh"] = 2 * rbh
        return self.treeQualities["dbh"]

    def get_end_nodes(self) -> list[int]:
        """
        Returns a list of end nodes in the graph.

        If the end nodes have already been computed, it returns the saved result.
        Otherwise, it computes the end nodes by finding nodes with a degree of 1 in the graph.

        Returns:
            A list of end nodes in the graph.
        """
        if not self.digraph:
            log.warning(
                "Graph not initialized, running initialize_graph_from with default values"
            )
            self.initialize_digraph_from()
        g = self.digraph

        if self.end_nodes:
            return self.end_nodes
        elif len(g.nodes) > 0:
            # the root node also has degree one, so we exclude it with n!= -1
            end_nodes = [n for n in g.nodes if g.degree(n) == 1 and n != -1]
            self.end_nodes = end_nodes
            return end_nodes

    def get_trunk_nodes(self) -> list[int]:
        """
        Returns a list of trunk nodes in the digraph.

        If the trunk nodes have already been computed, it returns the saved result.
        Otherwise, it computes the trunk nodes based on the branch order of the cylinders.

        Returns:
            A list of node ids for trunk nodes in the digraph.
        """
        g = self.digraph
        if self.trunk_nodes:
            return self.trunk_nodes
        elif len(g.nodes) > 0:
            trunk_cyls, _ = lam_filter(self.cylinders, lambda: branch_order == 0)  # noqa
            trunk_nodes = [cyl.cyl_id for cyl in trunk_cyls]
            self.trunk_nodes = trunk_nodes
            return trunk_nodes
        else:
            log.warning(
                "Graph not initialized, run <CylinderCollection>.initialize_graph(**args)"
            )
            return list(None)

    def direction_from_trunk(self, cyl: Cylinder):
        """
            Returns an array representing the direction of the cylinder from the trunk.

        Args:
            cyl: the cylinder for which to determine the direction from the trunk.

        Returns:
            tuple:
                A length 3 tuple corresponding the the x, y and z location of the cylinder.
                For each, a value of -1 indicates the cylinders location on the given axis
                    is less than the trunk's location on the same axis.
        """
        base_cyl = self.cylinders[0]
        tree_center = (base_cyl.x[0], base_cyl.y[0], base_cyl.z[0])
        x, y, z = cyl.x[1], cyl.y[1], cyl.z[1]
        directon_from_trunk = (
            -1 if x <= tree_center[0] else 1,
            -1 if y <= tree_center[1] else 1,
            -1 if z <= tree_center[1] else 1,
        )
        return directon_from_trunk

    def watershed_boundary(
        self,
        plane: str = "XY",
        filter_lambda: Callable | None = None,
        curvature_alpha: np.float64 = 1.8,
        draw: bool = False,
        save: bool = False,
        file_name_ext: str = "",
    ) -> None:
        """
        Generates tightly fit concave_hull (alpha shape) around the passed component.

        Args:
            component: The component for which to generate the watershed boundary. If None, the entire digraph is used.
            plane: The plane in which to generate the watershed boundary. Defaults to "XY".
            filter_lambda:
                A lambda function that determines which cylinders to include in the drawing.
                Takes no inputs but may reference any attribute of a cylinder object.
            curvature_alpha: The alpha value that determines the tightness of the fit of the shape.
                             Too low an alpha results in multiple shapes. Defaults to 2.
            draw: If True, the hull polygon is drawn via matplotlib. Defaults to False.
            save: If True, the hull is drawn and saved. Defaults to False.
            file_name_ext:
                The file extension to add to the end of the file name when when saving the drawn hull. Defaults to an empty string.

        Returns:
            None
        """
        comp_cylinders = self.cylinders
        if filter_lambda:
            comp_cylinders, _ = lam_filter(self.cylinders, filter_lambda)

        end_nodes = self.get_end_nodes()
        end_cylinders = [
            cyl
            for cyl in comp_cylinders
            if cyl.cyl_id in end_nodes or cyl.branch_order == 0
        ]
        boundary_points = [
            (cyl.x[1], cyl.y[1], cyl.z[1] + cyl.radius * 2) for cyl in comp_cylinders
        ]
        if plane in ("YZ", "XZ"):
            # using only branch tips as boundary points only really
            #  works in the XY direction. In the other directions
            #  the botom of the canopy is bounded by points on the surfaces
            # of various cylinders
            comp_cylinders = [cyl for cyl in comp_cylinders if cyl.branch_order != 0]
            for cyl in comp_cylinders:
                buffer = cyl.radius * 2
                buffer_arr = np.array(self.direction_from_trunk(cyl)) * buffer
                coords = (cyl.x[0], cyl.y[0], cyl.z[0])
                buffered_coords = coords + buffer_arr
                boundary_points.append(buffered_coords)

        if plane == "XY":
            boundary_points = [Point(cyl.x[1], cyl.y[1]) for cyl in end_cylinders]
        elif plane == "XZ":
            boundary_points = [Point(x, z) for x, _, z in boundary_points]
        elif plane == "YZ":
            boundary_points = [Point(y, z) for _, y, z in boundary_points]
        hull, _, _ = concave_hull(boundary_points, curvature_alpha)
        if draw:
            draw_cyls([hull], save=save, file_name_ext=file_name_ext)
        self.hulls[plane] = hull
        return hull, boundary_points

    def initialize_digraph_from(self, in_flow_grade_lim=in_flow_grade_lim):
        """This function creates a directed graph and its undirected counterpart.
        Initializes edge attributes as cylinder objects

        Args:
            in_flow_grade_lim: The drip edge cut off angle.
            Defaults to the in_flow_grade_lim from the configuration file.

        """
        gr = nx.DiGraph()
        trunk_nodes, _ = lam_filter(self.cylinders, lambda: branch_order == 0)  # noqa
        edges = (
            (
                (int(cyl.cyl_id), int(cyl.parent_id), {"cylinder": cyl})
                if (cyl.angle >= in_flow_grade_lim or cyl.cyl_id in trunk_nodes)
                else (int(cyl.parent_id), int(cyl.cyl_id), {"cylinder": cyl})
            )
            for cyl in self.cylinders
        )
        gr.add_edges_from(edges)
        self.digraph = gr
        self.graph = gr.to_directed()

    def contracted_nodes(self, node_list, next_node, G) -> tuple:
        """
        Returns a graph made from removing nodes from G,
        and connecting the removed nodes' neighbors via an edge

        Args:
            node_list:
                Nodes to be removed from the graph and connected to the next node.
            next_node:
                The node to which the removed nodes should be connected.
            G:
                The graph on which to carry out the operation

        Returns:
            nx.Graph: A graph with the listed nodes removed and their neighbors connected to the next node.
            list: A list of the neighbors of the removed nodes
        """
        # finding and removing the listed nodes to find their edges
        contracted = copy.deepcopy(G)
        contracted.remove_nodes_from(node_list)

        neighbors = [
            node
            for node in contracted.nodes()
            if (G.degree(node) != contracted.degree(node))
        ]

        contracted.add_node(next_node)
        edges_to_trunk = [(u, next_node) for u in neighbors]
        contracted.add_edges_from(edges_to_trunk)

        return contracted, neighbors

    def find_trunk_distances(self) -> dict:
        """
        Finds the distance in the graph (in number of nodes) between each node and the closest trunk node
        """
        trunk_nodes = self.get_trunk_nodes()

        trunk_contraction, _ = self.contracted_nodes(trunk_nodes, 0, self.graph)
        trunk_paths = nx.shortest_path(trunk_contraction, target=0)
        dists = {node: len(path) - 1 for node, path in trunk_paths.items()}
        return dists

    def find_flow_components(self, save_components: bool | None = False):
        """
        Determines to which partition perciplitation intercepted by
            each cylinder in the collection is likely to contribute to.
        That is, viewing each cylinder as a source, and allowing travel on the
            digraph as defined by the edge's direction, the function determines:
            - the set of nodes that can be reached from each source node.
            - for each set of nodes (each, 'flow') the eventual sink (or drip nodes)
                to which the theoretical flow is directed

        Args:
            save_components:
                If true, the stemflow and drip components of the collection's graph
                are saved. This is rather memory intensive so this input defaults to False.

        Raises:
            TypeError: Raised if the function is called on an undirected graph.
        """
        g = self.digraph
        if type(g) == nx.Graph:
            msg = "Find Flow Digraph invoked for undirected graph"
            log.error(msg)
            raise TypeError(msg)
        root_node = -1

        divide_nodes = [
            node
            for node in g.nodes()
            if g.out_degree(node) > 1
            or (g.out_degree(node) == 1 and g.in_degree(node) == 0)
        ]
        drip_nodes = [
            node
            for node, out_degree in g.out_degree()
            if out_degree == 0 and node != -1
        ]
        stem_flow_component = g.subgraph(nx.ancestors(g, root_node) | {0}).copy()

        stem_cylinders = [
            node
            for node in stem_flow_component.nodes()
            if stem_flow_component.degree(node) > 0
        ]
        for cyl in self.cylinders:
            if cyl.cyl_id in stem_cylinders:
                cyl.is_stem = True

        g_drip = copy.deepcopy(g)
        g_drip.remove_edges_from(stem_flow_component.edges())
        # g_drip.remove_nodes_from(stem_flow_component.nodes())

        drip_divide_pairings = [
            (
                drip_node,
                [node for node in divide_nodes if nx.has_path(g_drip, node, drip_node)],
            )
            for drip_node in drip_nodes
            if drip_node != -1
        ]

        drip_components = []
        cyl_to_drip_node = {cyl.cyl_id: [] for cyl in self.cylinders}
        for pair in drip_divide_pairings:
            drip_node, source_divides = pair
            paths = [
                nx.shortest_path(g_drip, div_node, drip_node)
                for div_node in source_divides
            ]
            component_nodes = chain.from_iterable(paths)
            component_graph = g_drip.subgraph(component_nodes).copy()
            component_cyls = [
                attr["cylinder"].cyl_id
                for _, _, attr in component_graph.edges(data=True)
            ]
            for idx in component_cyls:
                cyl_to_drip_node[idx].append(drip_node)

            drip_components.append((drip_node, g_drip.subgraph(component_nodes).copy()))

        for cyl in self.cylinders:
            if cyl_to_drip_node[cyl.cyl_id]:
                cyl.drip_node = cyl_to_drip_node[cyl.cyl_id][-1]

        log.info(
            f"{self.file_name} found to have {len(drip_components)} drip components"
        )
        print("reached_End of find flows")
        self.divide_nodes = divide_nodes
        self.drip_nodes = drip_nodes
        self.cyl_to_drip = cyl_to_drip_node
        if save_components:
            self.stem_flow_component = stem_flow_component
            self.drip_flow_components = drip_components
            self.drip_graph = g_drip

    # @profile
    def calculate_flows(self, plane: str = "XY"):
        """
        Based on the findings of 'find_flows', sums the qualities of
            all of the cylinders in each flow and stores the results in
            flow objects.
        Args:
            plane:
                The projection for which to source the angle and projected area
                data sumed. Valid values are "XY", "XZ", "YZ", and "3D".
        """
        if not self.projections[plane]:
            log.warning(
                f"""No {plane} projection data found. Running project_cylinders
                            This may take several minutes"""
            )
            self.project_cylinders(plane)
        cyls = self.cylinders
        flow_chars = []
        # edge_attributes = {}

        # log.info("attempting to sum stem edges ")

        # this probably doesn't belong here but its efficient to do it now
        # is 'needed' for statistics section
        np_flow_chars = [None] * (len(self.drip_nodes) + 1)

        def numpy_flow_chars(lambda_filter: Callable, drip_cyl, index: int):
            arr = np.array(
                [
                    np.array(
                        [
                            1,
                            np.float64(cyl.projected_data[plane]["area"]),
                            cyl.surface_area,
                            cyl.angle,
                            cyl.volume,
                            cyl.sa_to_vol,
                        ]
                    )
                    for cyl in cyls
                    if lambda_filter(cyl)
                ]
            )
            flow = np.sum(arr, axis=0)
            np_flow_chars[index] = Flow(
                flow[0],
                flow[1],
                flow[2],
                flow[3],
                flow[4],
                flow[5],
                drip_cyl.cyl_id,
                (drip_cyl.x[0], drip_cyl.y[0], drip_cyl.z[0]),
            )

        numpy_flow_chars(
            lambda_filter=lambda x: x.is_stem, drip_cyl=self.cylinders[0], index=0
        )
        log.info(f"np flow chars {np_flow_chars}")
        log.info(f" flow chars {flow_chars}")
        flow_chars.append(np_flow_chars[0])
        log.info(f" togetha {flow_chars}")
        # log.info(f"summed stem edges {flow_chars}")
        for idx, drip_node in enumerate(self.drip_nodes):
            cyl_before_drip = [cyl for cyl in cyls if cyl.cyl_id == drip_node]

            if len(cyl_before_drip) > 1:
                log.warning(f"Error: More that 1 cyl with id {drip_node} found")

            drip_node_loc = (
                cyl_before_drip[0].x[1],
                cyl_before_drip[0].y[1],
                cyl_before_drip[0].z[1],
            )
            drip_cylinders = [
                cyl
                for cyl in cyls
                if cyl.drip_node == drip_node and cyl.cyl_id != drip_node
            ]

            flow_chars.append(
                Flow(
                    **{
                        "num_cylinders": len(drip_cylinders),
                        "projected_area": np.sum(
                            [
                                np.float64(cyl.projected_data[plane]["area"])
                                for cyl in drip_cylinders
                            ]
                        ),
                        "surface_area": np.sum(
                            [cyl.surface_area for cyl in drip_cylinders]
                        ),
                        "angle_sum": np.sum([cyl.angle for cyl in drip_cylinders]),
                        "volume": np.sum([cyl.volume for cyl in drip_cylinders]),
                        "sa_to_vol": np.sum([cyl.sa_to_vol for cyl in drip_cylinders]),
                        "drip_node_id": drip_node,
                        "drip_node_loc": drip_node_loc,
                    }
                )
            )
        self.flows = flow_chars

    def identify_stem_paths(self):
        """
        Identifies the paths in the collection's graph by which the root node can be reached.
        """
        # Special case tree attributes
        root = 0
        self.stem_paths = nx.all_simple_paths(self.g, 0, root)

    def find_furthest(node, nodes):
        """Finds the furthest node in the given list from the provided n"""
        nodes = np.asarray(nodes)
        dist_2 = np.sum((nodes - node) ** 2, axis=1)
        return np.argmax(dist_2)

    def find_trunk_lean(self):
        """
        Defines a straight line from base to tip of the trunk then
         finds the angle of that line from the XZ plane
        """
        trunk_beg = lam_filter(self.cylinders, a_lambda=lambda: branch_order == 0)  # noqa
        trunk_points = [(cyl.x[0], cyl.y[0], cyl.z[0]) for cyl in trunk_beg]
        root = trunk_points[0]
        furthest_afeild, _ = furthest_point(root, trunk_points)
        dx, dy, dz = ((a - b) for a, b in zip(furthest_afeild, root))
        run = np.sqrt(dx**2 + dy**2)
        angle = np.arctan(dz / np.sqrt(dx**2 + dy**2)) if run > 0 else np.arctan(0)
        self.trunk_lean = angle
        return angle

    def find_overlap_by_percentile(
        self, plane: str = "XY", percentiles: list[int] = [25, 50, 75]
    ):
        """
        Groups the projected polygons of the collection's cylinders
            by percentiles given - with respect to their z values.
        Passes the list of grouped polygons to the get_projected_overlap function
            to get the overlap percentages between the groups.
        See get_projected_overlap for more information.

        Args:
            plane:
                Plane of projection from which to source the polygons.
                Defaults to "XY".
            percentiles:
                Percentiles defining the groups passed to get_projected_overlap.
                Expected to be strictly increasing. Defaults to [25, 50, 75].

        Returns:
            list[dict]: Each dicitonary corresponding to a percentile,
            with metrics describing the overlap between the polygons in that percentile
            {
                "sum_area": sum of areas of the polygons in the percentile,
                "effective_area": the area of the union of the polygons in the percentile,
                "internal_overlap": sum_area - effective_area,
                "overlap_with_previous": the area of the overlap between the polygons in the percentile,
                        with the polygons in the previous percentile,
            }
        """
        percentiles.sort()
        if not self.projections[plane]:
            self.project_cylinders(plane)
        non_trunk_polys, _ = lam_filter(self.cylinders, lambda: branch_order != 0)  # noqa
        cyl_metric = [cyl.z[0] for cyl in non_trunk_polys]
        percentiles_by_metric = np.percentile(cyl_metric, percentiles)
        prev_perc = 0
        poly_matrix = []
        for perc in percentiles_by_metric:
            poly_matrix.append(
                [
                    cyl.projected_data[plane]["polygon"]
                    for cyl in non_trunk_polys
                    if cyl.z[0] <= perc and cyl.z[0] > prev_perc
                ]
            )
            prev_perc = perc
        overlaps = get_projected_overlap(poly_matrix, percentiles)

        return overlaps

    def statistics(self, plane: str = "XY", file_name_ext: str = ""):
        """Generates a csv and dictionary containing a variety of descriptive
            statistics about the collection of cylinders.

        Args:
            plane:
                The plane used for projection dependent statistics.
                Defaults to "XY".
            file_name_ext:
                the naming suffix to used for the saved file name

        Returns:
            _description_
        """
        if not self.pSV:
            self.project_cylinders(plane)
        if not self.hulls[plane]:
            self.watershed_boundary(plane=plane)
        if not self.stem_hull:
            if not self.stem_flow_component:
                self.find_flow_components()
                self.calculate_flows(plane=plane)
            self.stem_hull, _ = self.watershed_boundary(
                plane="XY",
                filter_lambda=lambda: is_stem,  # noqa
            )
        dbh = self.get_dbh()

        canopy_cover = self.hulls[plane].area
        canopy_boundary = self.hulls[plane].boundary.length

        canopy_cover_stem = self.stem_hull.area
        canopy_boundary_stem = self.stem_hull.boundary.length

        log.info("Found hull alpha shape stats")

        # calculate projected areas and (therefore) overlaps
        # Unary union gives a single contiguous polygon when fed many overlapping cylinders
        # The area of the union thus differs from the sum of the areas of its components
        # in that the former counts overlaps only once
        # polys = [poly_dict['polygon'] for poly_dict in self.pSV]
        polys = self.pSV
        tot_poly = unary_union(polys)
        projected_union_area = tot_poly.area
        sum_projected_area = np.sum([poly.area for poly in self.pSV])
        stem_polys = [
            cyl.projected_data[plane]["polygon"]
            for cyl in self.cylinders
            if cyl.is_stem
        ]
        union_poly_stem = unary_union(stem_polys)
        projected_union_area_stem = union_poly_stem.area
        sum_projected_area_stem = np.sum([poly.area for poly in stem_polys])

        log.info("found projected areas")

        # this could techically be conmbined with the above by adding a percentile of 100
        #
        overlap_dict = self.find_overlap_by_percentile(percentiles=[25, 50, 75])

        tot_poly = unary_union(self.pSV)

        min_x = self.extent["min"][0]
        max_x = self.extent["max"][0]
        min_y = self.extent["min"][1]
        max_y = self.extent["max"][1]
        min_z = self.extent["min"][2]
        max_z = self.extent["max"][2]

        if self.flows:
            stem_flow = [flow for flow in self.flows if flow.drip_node_id == 0][0]
        else:
            stem_flow = 0

        total_surface_area = np.sum([cyl.surface_area for cyl in self.cylinders])
        total_volume = np.sum([cyl.volume for cyl in self.cylinders])
        max_bo = np.max([cyl.branch_order for cyl in self.cylinders])

        order_zero_cyls, _ = lam_filter(self.cylinders, lambda: branch_order == 0)  # noqa
        order_one_cyls, _ = lam_filter(self.cylinders, lambda: branch_order == 1)  # noqa
        order_two_cyls, _ = lam_filter(self.cylinders, lambda: branch_order == 2)  # noqa
        order_three_cyls, _ = lam_filter(self.cylinders, lambda: branch_order == 3)  # noqa
        order_four_cyls, _ = lam_filter(self.cylinders, lambda: branch_order == 4)  # noqa

        statistics = {
            "total_psa": projected_union_area,
            "psa_w_overlap": sum_projected_area,
            "stem_psa": projected_union_area_stem,
            "stem_psa_w_overlap": sum_projected_area_stem,
            "tot_surface_area": total_surface_area,
            "stem_surface_area": stem_flow.surface_area,
            "tot_hull_area": canopy_cover,
            "tot_hull_boundary": canopy_boundary,
            "stem_hull_area": canopy_cover_stem,
            "stem_hull_boundary": canopy_boundary_stem,
            "num_drip_points": len(self.drip_nodes),
            "max_bo": max_bo,
            "topQuarterTotPsa": overlap_dict[75]["sum_area"],
            "topHalfTotPsa": overlap_dict[50]["sum_area"],
            "topThreeQuarterTotPsa": overlap_dict[25]["sum_area"],
            "TotalShade": sum_projected_area - projected_union_area,
            "top_quarter_shade": overlap_dict[75]["overlap_with_previous"],
            "top_half_shade": overlap_dict[50]["overlap_with_previous"],
            "top_three_quarter_shade": overlap_dict[25]["overlap_with_previous"],
            "DBH": dbh,
            "volume": total_volume,
            "X_max": max_x,
            "Y_max": max_y,
            "Z_max": max_z,
            "X_min": min_x,
            "Y_min": min_y,
            "Z_min": min_z,
            "Order_zero_angle_avg": np.average([cyl.angle for cyl in order_zero_cyls]),
            "Order_zero_angle_std": np.std([cyl.angle for cyl in order_zero_cyls]),
            "Order_one_angle_avg": np.average([cyl.angle for cyl in order_one_cyls]),
            "Order_one_angle_std": np.std([cyl.angle for cyl in order_one_cyls]),
            "Order_two_angle_avg": np.average([cyl.angle for cyl in order_two_cyls]),
            "Order_two_angle_std": np.std([cyl.angle for cyl in order_two_cyls]),
            "Order_three_angle_avg": np.average(
                [cyl.angle for cyl in order_three_cyls]
            ),
            "Order_three_angle_std": np.std([cyl.angle for cyl in order_three_cyls]),
            "order_gr_four_angle_avg": np.average(
                [cyl.angle for cyl in order_four_cyls]
            ),
            "order_gr_four_angle_std": np.std([cyl.angle for cyl in order_four_cyls]),
            "file_name": self.file_name + plane,
        }
        stat_file = save_file(
            self.file_name.replace(".csv", f"_{file_name_ext}"),
            out_file=statistics,
            method="statistics",
            overwrite=True,
        )

        return stat_file

    def generate_flow_file(self, file_name_ext):
        """Creates a csv and dictionary containing data describing
        the flows calculated in find_flow_components and calculate_flows

        Args:
            file_name_ext:
                the naming suffix to used for the saved file name

        Returns:
            dict: A dictionary containing the flow data
        """
        flow_dicts = [flow.__dict__ for flow in self.flows]
        flow_file = save_file(
            self.file_name.replace(".csv", f"_flows_{file_name_ext}"),
            out_file=flow_dicts,
            subdir="flows",
            method="flows",
            overwrite=True,
        )
        return flow_file

    def set_drip_points(
        self,
        # metric :str = 'projected_area',
        percentile: int = 90,
    ):
        """
        Returns the locations of the drip points.
        **Drip points are identified as the subset of drip nodes
            with a value for the provided metric in the percentile given
        """
        if not self.flows:
            log.warning(
                """flows are not defined, running 'find_flow_components'
                            and  'calculate flows'"""
            )
            log.warning("This may take several minutes")
            self.find_flow_components()
            self.calculate_flows()
        # excluding trunk node
        drip_nodes = [
            (f.drip_node_id, f.projected_area, f.drip_node_loc)
            for f in self.flows
            if f.drip_node_id != 0
        ]

        distinct_drip_node_ids = {node[0] for node in drip_nodes}

        flow_by_area_dict = {}
        for node in distinct_drip_node_ids:
            area = np.sum([flow[1] for flow in drip_nodes if flow[0] == node])
            flow_by_area_dict[node] = area

        point_cutoff = np.percentile(
            [area for _, area in flow_by_area_dict.items()], percentile
        )
        drip_points = [
            node for node, area in flow_by_area_dict.items() if area > point_cutoff
        ]
        if drip_points == []:
            log.warning(f"No drip points found in the {percentile}th percentile")
            log.warning("Assuming a small drip point list, including all drip points")
            drip_points = [node for node, _ in flow_by_area_dict.items()]
        drip_point_locs = list(
            [node[2][0], node[2][1], node[2][2]]
            for node in drip_nodes
            if node[0] in drip_points
        )
        self.drip_point_loc = drip_point_locs
        return drip_point_locs
