import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
import matplotlib.colors as mcolors

from typing import Callable, List
from neuron import h

from dendrotweaks.morphology.trees import Node, Tree
from dendrotweaks.morphology.domains import Domain
from dataclasses import dataclass, field
from bisect import bisect_left

import warnings

from dendrotweaks.utils import get_domain_color

def custom_warning_formatter(message, category, filename, lineno, file=None, line=None):
    return f"{category.__name__}: {message} ({os.path.basename(filename)}, line {lineno})\n"

warnings.formatwarning = custom_warning_formatter



class Section(Node):
    """
    A class representing a section in a neuron morphology.

    A section is continuous part of a neuron's morphology between two branching points.

    Parameters
    ----------
    idx : str
        The index of the section.
    parent_idx : str
        The index of the parent section.
    points : List[Point]
        The points that define the section.

    Attributes
    ----------
    points : List[Point]
        The points that define the section.
    segments : List[Segment]
        The segments to which the section is divided.
    _ref : h.Section
        The reference to the NEURON section.
    """

    def __init__(self, idx: str, parent_idx: str, points: List[Node]) -> None:
        super().__init__(idx, parent_idx)
        self.domain_idx = None
        self.points = points
        self.segments = []
        self._ref = None
        self._domain = self.points[0].domain

        if not all(pt.domain == self._domain for pt in points):
            raise ValueError('All points in a section must belong to the same domain.')


    # MAGIC METHODS

    def __call__(self, x: float):
        """
        Return the segment at a given position.
        """
        if self._ref is None:
            raise ValueError('Section is not referenced in NEURON.')
        if x < 0 or x > 1:
            raise ValueError('Location x must be in the range [0, 1].')
        elif x == 0: 
            # TODO: Decide how to handle sec(0) and sec(1)
            # as they are not shown in the seg_graph
            return self.segments[0]
        elif x == 1:
            return self.segments[-1]
        matching_segs = [self._ref(x) == seg._ref for seg in self.segments]
        if any(matching_segs):
            return self.segments[matching_segs.index(True)]
        raise ValueError(f'No segment found at location {x}')
        

    def __iter__(self):
        """
        Iterate over the segments in the section.
        """
        for seg in self.segments:
            yield seg


    # PROPERTIES

    @property
    def domain(self):
        """
        The morphological or functional domain of the node.
        """
        return self._domain


    @domain.setter
    def domain(self, domain):
        self._domain = domain
        for pt in self.points:
            pt.domain = domain


    @property
    def df_points(self):
        """
        A DataFrame of the points in the section.
        """
        # concatenate the dataframes of the nodes
        return pd.concat([pt.df for pt in self.points])


    @property
    def df(self):
        """
        A DataFrame of the section.
        """
        return pd.DataFrame({'idx': [self.idx],
                            'parent_idx': [self.parent_idx]})


    # TODO: Figure out why this is different from NEURON's diam
    # Update: In NEURON, sec.diam returns the diameter of the segment at the center of the section
    # In this implementation, sec.diam returns the average diameter of the section
    # @property
    # def diam(self):
    #     """
    #     Average diameter of the section calculated from
    #     the radii and distances of the points.
    #     """
    #     distances = self.distances  # Cumulative distances
    #     radii = self.radii  # Corresponding radii
    #     total_length = distances[-1]  # Total section length
        
    #     if total_length == 0:
    #         return 0  # Avoid division by zero for zero-length sections

    #     segment_lengths = np.diff(distances)  # Lengths of frusta segments
    #     segment_diameters = 2 * (np.array(radii[:-1]) + np.array(radii[1:])) / 2  # Mean diameter per segment

    #     # Length-weighted average
    #     avg_diameter = np.sum(segment_diameters * segment_lengths) / total_length

    #     return avg_diameter

    @property
    def diam(self):
        """
        Diameter of the central segment of the section (from NEURON).
        """
        return self._ref.diam
        

    @property
    def L(self):
        """
        Length of the section (from NEURON).
        """
        return self._ref.L


    @property
    def cm(self):
        """
        Specific membrane capacitance of the section (from NEURON).
        """
        return self._ref.cm


    @property
    def Ra(self):
        """
        Axial resistance of the section (from NEURON).
        """
        return self._ref.Ra


    @property
    def nseg(self):
        """
        Number of segments in the section (from NEURON).
        """
        return self._ref.nseg

    @nseg.setter
    def nseg(self, value):
        if value < 1:
            raise ValueError('Number of segments must be at least 1.')
        if value % 2 == 0:
            raise ValueError('Number of segments must be odd.')
        # Set the number in NEURON
        self._ref.nseg = value
        # Get the new NEURON segments
        nrnsegs = [seg for seg in self._ref]

        # Create new DendroTweaks segments
        from dendrotweaks.morphology.seg_trees import Segment
        old_segments = self.segments
        new_segments = [Segment(idx=0, parent_idx=0, neuron_seg=seg, section=self)
                            for seg in nrnsegs]
        
        seg_tree = self._tree._seg_tree
        first_segment = self.segments[0]
        parent = first_segment.parent

        for i, seg in enumerate(new_segments[:]):
            if i == 0:
                seg_tree.insert_node_before(seg, first_segment)
            else:
                seg_tree.insert_node_before(seg, new_segments[i-1])

        for seg in old_segments:
            seg_tree.remove_node(seg)
        
        # Sort the tree
        self._tree._seg_tree.sort()
        # Update the section's segments
        self.segments = new_segments


    @property
    def radii(self):
        """
        Radii of the points in the section.
        """
        return [pt.r for pt in self.points]


    @property
    def diameters(self):
        """
        Diameters of the points in the section.
        """
        return [2 * pt.r for pt in self.points]


    @property
    def xs (self):
        """
        X-coordinates of the points in the section.
        """
        return [pt.x for pt in self.points]


    @property
    def ys(self):
        """
        Y-coordinates of the points in the section.
        """
        return [pt.y for pt in self.points]


    @property
    def zs(self):
        """
        Z-coordinates of the points in the section.
        """
        return [pt.z for pt in self.points]


    @property
    def seg_centers(self):
        """
        The list of segment centers in the section with normalized length.
        """
        if self._ref is None:
            raise ValueError('Section is not referenced in NEURON.')
        return (np.array([(2*i - 1) / (2 * self._ref.nseg)
                for i in range(1, self._ref.nseg + 1)]) * self._ref.L).tolist()


    @property
    def seg_borders(self):
        """
        The list of segment borders in the section with normalized length.
        """
        if self._ref is None:
            raise ValueError('Section is not referenced in NEURON.')
        nseg = int(self._ref.nseg)
        return [i / nseg for i in range(nseg + 1)]


    @property
    def distances(self):
        """
        The list of cumulative euclidean distances of the points in the section.
        """
        coords = np.array([[pt.x, pt.y, pt.z] for pt in self.points])
        deltas = np.diff(coords, axis=0)
        frusta_distances = np.sqrt(np.sum(deltas**2, axis=1))
        cumulative_frusta_distances = np.insert(np.cumsum(frusta_distances), 0, 0)
        return cumulative_frusta_distances


    @property
    def center(self):
        """
        The coordinates of the center of the section.
        """
        return np.mean(self.xs), np.mean(self.ys), np.mean(self.zs)


    @property
    def length(self):
        """
        The length of the section calculated as the sum of the distances between the points.
        """
        return self.distances[-1]


    @property
    def area(self):
        """
        The surface area of the section calculated as the sum of the areas of the frusta segments.
        """
        areas = [np.pi * (r1 + r2) * np.sqrt((r1 - r2)**2 + h**2) for r1, r2, h in zip(self.radii[:-1], self.radii[1:], np.diff(self.distances))]
        return sum(areas)

    def has_mechanism(mech_name):
        """
        Check if the section has a mechanism inserted.

        Parameters
        ----------
        mech_name : str
            The name of the mechanism to check.
        """
        return self._ref.has_membrane(mech_name)


    # REFERENCING METHODS

    def create_and_reference(self, simulator_name='NEURON'):
        """
        Add a reference to the section in the simulator.

        Parameters
        ----------
        simulator_name : str
            The name of the simulator to create the section in.
        """
        if simulator_name == 'NEURON':
            self.create_NEURON_section()
        elif simulator_name == 'JAXLEY':
            self.create_JAXLEY_section()


    def create_NEURON_section(self):
        """
        Create a NEURON section.
        """
        self._ref = h.Section() # name=f'Sec_{self.idx}'
        if self.parent is not None:
            # TODO: Attaching basal to soma 0
            if self.parent.parent is None: # if parent is soma
                self._ref.connect(self.parent._ref(0.5))
            else:
                self._ref.connect(self.parent._ref(1))
        # Add 3D points to the section
        self._ref.pt3dclear()
        for pt in self.points:
            diam = 2*pt.r
            diam = round(diam, 16)
            self._ref.pt3dadd(pt.x, pt.y, pt.z, diam)

    def create_JAXLEY_section(self):
        """
        Create a JAXLEY section.
        """
        raise NotImplementedError


    # MECHANISM METHODS

    def insert_mechanism(self, name: str):
        """
        Inserts a mechanism in the section if 
        it is not already inserted.
        """
        if self._ref.has_membrane(name):
            return
        self._ref.insert(name)


    def uninsert_mechanism(self, name: str):
        """
        Uninserts a mechanism in the section if
        it was inserted.
        """
        if not self._ref.has_membrane(name):
            return
        self._ref.uninsert(name)


    # PARAMETER METHODS

    def get_param_value(self, param_name):
        """
        Get the average parameter of the section's segments.

        Parameters
        ----------
        param_name : str
            The name of the parameter to get.

        Returns
        -------
        float
            The average value of the parameter in the section's segments.
        """
        # if param_name in ['Ra', 'diam', 'L', 'nseg', 'domain', 'subtree_size']:
        #     return getattr(self, param_name)
        # if param_name in ['dist']:
        #     return self.distance_to_root(0.5)
        seg_values = [seg.get_param_value(param_name) for seg in self.segments]
        return round(np.mean(seg_values), 16)


    def path_distance(self, relative_position: float = 0, 
                        within_domain: bool = False) -> float:
        """
        Calculate the distance from the section to the root at a given relative position.

        Parameters
        ----------
        relative_position : float
            The position along the section's normalized length [0, 1].
        within_domain : bool
            Whether to stop at the domain boundary.

        Returns
        -------
        float
            The distance from the section to the root.

        Important
        ---------
        Assumes that we always attach the 0 end of the child.
        """
        if not (0 <= relative_position <= 1):
            raise ValueError('Relative position must be between 0 and 1.')

        distance = 0
        factor = relative_position
        node = self

        while node.parent:

            distance += factor * node.length

            if within_domain and node.parent.domain != node.domain:
                break

            node = node.parent
            factor = 1
            
        return distance       

    
    def disconnect_from_parent(self):
        """
        Detach the section from its parent section.
        """
        # In SectionTree
        super().disconnect_from_parent()
        # In NEURON
        if self._ref:
            h.disconnect(sec=self._ref) #from parent
        # In PointTree
        self.points[0].disconnect_from_parent()
        # In SegmentTree
        if self.segments:
            self.segments[0].disconnect_from_parent()


    def connect_to_parent(self, parent):
        """
        Attach the section to a parent section.

        Parameters
        ----------
        parent : Section
            The parent section to attach to.
        """
        # In SectionTree
        super().connect_to_parent(parent)
        # In NEURON
        if self._ref:
            if self.parent is not None:
                if self.parent.parent is None: # if parent is soma
                    self._ref.connect(self.parent._ref(0.5))
                else:
                    self._ref.connect(self.parent._ref(1))
                
        # In PointTree
        if self.parent is not None:
            if self.parent.parent is None: # if parent is soma
                parent_sec = self.parent
                parent_pt = parent_sec.points[1] if len(parent_sec.points) > 1 else parent_sec.points[0]
                self.points[0].connect_to_parent(parent_pt) # attach to the middle of the parent
            else:
                self.points[0].connect_to_parent(parent.points[-1]) # attach to the end of the parent
        # In SegmentTree
        if self.segments:
            self.segments[0].connect_to_parent(parent.segments[-1])


    # PLOTTING METHODS
   
    def plot(self, ax=None, plot_parent=True, section_color=None, parent_color='gray', 
             show_labels=True, aspect_equal=True):
        """
        Plot section morphology in 3D projections (XZ, YZ, XY) and radii distribution.
        
        Parameters
        ----------
        ax : list or array of matplotlib.axes.Axes, optional
            Four axes for plotting (XZ, YZ, XY, radii). If None, creates a new figure with axes.
        plot_parent : bool, optional
            Whether to include parent section in the visualization.
        section_color : str or None, optional
            Color for the current section. If None, assigns based on section domain.
        parent_color : str, optional
            Color for the parent section.
        show_labels : bool, optional
            Whether to show axis labels and titles.
        aspect_equal : bool, optional
            Whether to set aspect ratio to 'equal' for the projections.
            
        Returns
        -------
        ax : list of matplotlib.axes.Axes
            The axes containing the plots.
        """
        # Create figure and axes if not provided
        if ax is None:
            fig = plt.figure(figsize=(10, 8))
            gs = GridSpec(2, 3, width_ratios=[1, 1, 1.2], figure=fig)
            
            # Create the three projection axes and one radius axis
            ax_xz = fig.add_subplot(gs[0, 0])
            ax_yz = fig.add_subplot(gs[0, 1])
            ax_xy = fig.add_subplot(gs[1, 0])
            ax_radii = fig.add_subplot(gs[1, 1:])
            ax = [ax_xz, ax_yz, ax_xy, ax_radii]
        else:
            # Use provided axes
            if len(ax) != 4:
                # flatten 
                ax = [ai for a in ax for ai in a]
            ax_xz, ax_yz, ax_xy, ax_radii = ax
        
        # Determine section color based on domain if not provided
        if section_color is None:
            section_color = get_domain_color(self.domain)
        
        # Extract coordinates
        xs = np.array([p.x for p in self.points])
        ys = np.array([p.y for p in self.points])
        zs = np.array([p.z for p in self.points])
        
        # Plot section projections
        self._plot_projection(ax_xz, xs, zs, 'X', 'Z', 'XZ Projection', 
                              section_color, show_labels, aspect_equal)
        self._plot_projection(ax_yz, ys, zs, 'Y', 'Z', 'YZ Projection', 
                              section_color, show_labels, aspect_equal)
        self._plot_projection(ax_xy, xs, ys, 'X', 'Y', 'XY Projection', 
                              section_color, show_labels, aspect_equal)
        
        # Plot radius distribution
        self._plot_radii_distribution(ax_radii, plot_parent, section_color, parent_color)
        
        # Plot parent section if requested
        if plot_parent and self.parent:
            # Only plot parent projections, radii are handled in _plot_radii_distribution
            parent_xs = np.array([p.x for p in self.parent.points])
            parent_ys = np.array([p.y for p in self.parent.points])
            parent_zs = np.array([p.z for p in self.parent.points])
            
            self.parent._plot_projection(ax_xz, parent_xs, parent_zs, None, None, None, 
                                         parent_color, False, aspect_equal)
            self.parent._plot_projection(ax_yz, parent_ys, parent_zs, None, None, None, 
                                         parent_color, False, aspect_equal)
            self.parent._plot_projection(ax_xy, parent_xs, parent_ys, None, None, None, 
                                         parent_color, False, aspect_equal)
        
        # Add overall title if we created the figure
        if ax is not None and show_labels:
            fig = ax_xz.get_figure()
            fig.suptitle(f"Section {self.idx} ({self.domain})", fontsize=14)
            fig.tight_layout()
        
        return ax
    
    def _plot_projection(self, ax, x_coords, y_coords, x_label, y_label, title, 
                         color, show_labels, aspect_equal):
        """Helper method to plot a 2D projection of the section."""
        ax.plot(x_coords, y_coords, 'o-', color=color, markerfacecolor=color, 
                markeredgecolor='black', markersize=4, linewidth=1.5)
        
        if show_labels:
            if x_label:
                ax.set_xlabel(x_label)
            if y_label:
                ax.set_ylabel(y_label)
            if title:
                ax.set_title(title)
        
        if aspect_equal:
            ax.set_aspect('equal')
    
    def _plot_radii_distribution(self, ax, plot_parent, section_color, parent_color):
        """Helper method to plot radius distribution along the section."""
        # Get section length for normalization
        section_length = self.distances[-1] - self.distances[0]
        
        # Normalize distances to start at 0 and end at section_length
        normalized_distances = self.distances - self.distances[0]
        
        # Plot section radii
        ax.plot(normalized_distances, self.radii, 'o-', color=section_color, 
                label=f"{self.domain} ({self.idx})", linewidth=2)
        
        # Plot reference NEURON segments if available
        if hasattr(self, '_ref') and self._ref:
            # Calculate normalized segment centers
            normalized_seg_centers = np.array(self.seg_centers) - self.distances[0]
            
            # Extract radii from segments
            seg_radii = np.array([seg.diam / 2 for seg in self._ref])
            
            # Use the specified bar width calculation from original code
            bar_width = [self._ref.L / self._ref.nseg] * self._ref.nseg
            
            # Plot segment radii as bars
            ax.bar(normalized_seg_centers, seg_radii, width=bar_width, 
                   alpha=0.5, color=section_color, edgecolor='white',
                   label=f"{self.domain} segments")
        
        # Plot parent section if requested
        if plot_parent and self.parent:
            parent_length = self.parent.distances[-1] - self.parent.distances[0]
            
            # Normalize parent distances to end at 0 (connecting to child)
            # Parent section goes from -parent_length to 0
            normalized_parent_distances = self.parent.distances - self.parent.distances[-1]
            
            # Plot parent radii
            ax.plot(normalized_parent_distances, self.parent.radii, 'o-', 
                    color=parent_color, linewidth=2,
                    label=f"Parent {self.parent.domain} ({self.parent.idx})")
            
            # Plot parent reference segments if available
            if hasattr(self.parent, '_ref') and self.parent._ref:
                # Normalize parent segment centers to the same scale
                normalized_parent_seg_centers = (np.array(self.parent.seg_centers) - 
                                                self.parent.distances[-1])
                
                # Extract parent segment radii
                parent_seg_radii = np.array([seg.diam / 2 for seg in self.parent._ref])
                
                # Use the specified bar width calculation for parent
                parent_bar_width = [self.parent._ref.L / self.parent._ref.nseg] * self.parent._ref.nseg
                
                # Plot parent segment radii as bars
                ax.bar(normalized_parent_seg_centers, parent_seg_radii, 
                       width=parent_bar_width, alpha=0.5, color=parent_color, 
                       edgecolor='white', label=f"Parent segments")
        
        # Set plot labels and legend
        ax.set_xlabel('Distance (µm)')
        ax.set_ylabel('Radius (µm)')
        ax.set_title('Radius Distribution')
        
        # Ensure y-axis starts at 0
        ax.set_ylim(bottom=0)
        
        # Adjust x-axis to show the full section(s)
        if plot_parent and self.parent:
            parent_length = self.parent.distances[-1] - self.parent.distances[0]
            ax.set_xlim(-parent_length * 1.05, section_length * 1.05)
        else:
            ax.set_xlim(-section_length * 0.05, section_length * 1.05)
        
        # Add legend if we have multiple data series
        if ((hasattr(self, '_ref') and self._ref) or 
            (plot_parent and self.parent)):
            ax.legend(loc='best', frameon=True, framealpha=0.8)

    def plot_radii(self, ax=None, include_parent=False, section_color=None, parent_color='gray'):
        """
        Plot just the radius distribution for the section.
        
        Parameters
        ----------
        ax : matplotlib.axes.Axes, optional
            Axes to plot on. If None, creates a new figure and axes.
        include_parent : bool, optional
            Whether to include parent section in the plot.
        section_color : str or None, optional
            Color for current section. If None, assigns based on section domain.
        parent_color : str, optional
            Color for parent section if included.
            
        Returns
        -------
        ax : matplotlib.axes.Axes
            The axes containing the plot.
        """
        # Create new figure and axes if not provided
        if ax is None:
            fig, ax = plt.subplots(figsize=(10, 5))
        
        # Determine section color based on domain if not provided
        if section_color is None:
            domain_colors = {
                'soma': 'black',
                'axon': 'red',
                'dend': 'blue',
                'apic': 'green'
            }
            section_color = domain_colors.get(self.domain, 'purple')
        
        # Plot radius distribution
        self._plot_radii_distribution(ax, include_parent, section_color, parent_color)
        
        # Add title if creating a standalone plot
        if ax.get_figure().get_axes()[0] == ax:  # If this is the only axes in the figure
            ax.set_title(f"Radius Distribution - Section {self.idx} ({self.domain})")
            plt.tight_layout()
        
        return ax


class SectionTree(Tree):
    """
    A class representing a tree graph of sections in a neuron morphology.

    Parameters
    ----------
    sections : List[Section]
        A list of sections in the tree.

    Attributes
    ----------
    domains : Dict[str, Domain]
        A dictionary of domains in the tree.
    """

    def __init__(self, sections: list[Section]) -> None:
        super().__init__(sections)
        self._create_domains()
        self._point_tree = None
        self._seg_tree = None


    def __repr__(self):
        return f"SectionTree(root={self.root!r}, num_nodes={len(self._nodes)})"


    def _create_domains(self):
        """
        Create domains using the data from the sections (from the points in the sections).
        """

        unique_domain_names = set([sec.domain for sec in self.sections])
        self.domains = {name: Domain(name) for name in sorted(unique_domain_names)}

        for sec in self.sections:
            self.domains[sec.domain].add_section(sec)


    # PROPERTIES    

    @property
    def sections(self):
        """
        A list of sections in the tree. Alias for self._nodes.
        """
        return self._nodes


    @property
    def soma(self):
        """
        The soma section of the tree. Alias for self.root.
        """
        return self.root


    @property
    def sections_by_depth(self):
        """
        A dictionary of sections grouped by depth in the tree
        (depth is the number of edges from the root).
        """
        sections_by_depth = {}
        for sec in self.sections:
            if sec.depth not in sections_by_depth:
                sections_by_depth[sec.depth] = []
            sections_by_depth[sec.depth].append(sec)
        return sections_by_depth


    @property
    def df(self):
        """
        A DataFrame of the sections in the tree.
        """
        data = {
            'idx': [],
            'domain': [],
            'x': [],
            'y': [],
            'z': [],
            'r': [],
            'parent_idx': [],
            'section_idx': [],
            'parent_section_idx': [],
        }

        for sec in self.sections:
            points = sec.points if sec.parent is None or sec.parent.parent is None else sec.points[1:]
            for pt in points:
                data['idx'].append(pt.idx)
                data['domain'].append(pt.domain)
                data['x'].append(pt.x)
                data['y'].append(pt.y)
                data['z'].append(pt.z)
                data['r'].append(pt.r)
                data['parent_idx'].append(pt.parent_idx)
                data['section_idx'].append(sec.idx)
                data['parent_section_idx'].append(sec.parent_idx)

        return pd.DataFrame(data)


    def sort(self, **kwargs):
        """
        Sort the sections in the tree using a depth-first traversal.

        Parameters
        ----------
        sort_children : bool, optional
            Whether to sort the children of each node 
            based on the number of bifurcations in their subtrees. Defaults to True.
        force : bool, optional
            Whether to force the sorting of the tree even if it is already sorted. Defaults to False.
        """
        super().sort(**kwargs)
        self._point_tree.sort(**kwargs)
        if self._seg_tree:
            self._seg_tree.sort(**kwargs)


    # STRUCTURE METHODS

    def remove_subtree(self, section):
        """
        Remove a section and its subtree from the tree.

        Parameters
        ----------
        section : Section
            The section to remove.
        """
        super().remove_subtree(section)
        # Domains
        for domain in self.domains.values():
            for sec in section.subtree:
                if sec in domain.sections:
                    domain.remove_section(sec)
        # Points
        self._point_tree.remove_subtree(section.points[0])
        # Segments
        if self._seg_tree:
            self._seg_tree.remove_subtree(section.segments[0])
        # NEURON
        if section._ref:
            h.disconnect(sec=section._ref)
            for sec in section.subtree:
                h.delete_section(sec=sec._ref)


    def remove_zero_length_sections(self):
        """
        Remove sections with zero length.
        """
        for sec in self.sections:
            if sec.length == 0:
                for pt in sec.points:
                    self._point_tree.remove_node(pt)
                for seg in sec.segments:
                    self._seg_tree.remove_node(seg)
                self.remove_node(sec)


    def downsample(self, factor: float):
        """
        Downsample the SWC tree by reducing the number of points in each section 
        based on the given factor, while preserving the first and last points.
        
        :param factor: The proportion of points to keep (e.g., 0.5 keeps 50% of points)
                       If factor is 0, keep only the first and last points.
        """
        for sec in self.sections:
            if sec is self.soma:
                continue

            if len(sec.points) < 3:  # Keep sections with only start & end points
                continue

            num_points = len(sec.points)
            if factor == 0:
                num_to_keep = 2
            else:
                num_to_keep = max(2, int(num_points * factor))  # Ensure at least start & end remain

            # Select indices to keep (first, last, and spaced indices in between)
            keep_indices = np.linspace(0, num_points - 1, num_to_keep, dtype=int)
            keep_set = set(keep_indices)

            points_to_remove = [pt for i, pt in enumerate(sec.points) if i not in keep_set]

            print(f'Removing {len(points_to_remove)} points from section {sec.idx}')
            
            for pt in points_to_remove:
                self._point_tree.remove_node(pt)
                sec.points.remove(pt)
            
        self._point_tree.sort()

            
    # def plot_sections_as_matrix(self, ax=None):
    #     """
    #     Plot the sections as a connectivity matrix.
    #     """
    #     if ax is None:
    #         fig, ax = plt.subplots()

    #     n = len(self.sections)
    #     matrix = np.zeros((n, n))
    #     for section in self.sections:
    #         if section.parent:
    #             matrix[section.idx, section.parent.idx] = section.idx
    #     matrix[matrix == 0] = np.nan

    #     ax.imshow(matrix.T, cmap='jet_r')
    #     ax.set_xlabel('Section ID')
    #     ax.set_ylabel('Parent ID')


    # PLOTTING METHODS

    def plot(self, ax=None, show_points=False, show_lines=True, 
            show_domains=True, annotate=False, 
            projection='XY', highlight_sections=None, focus_sections=None):
        """
        Plot the sections in the tree in a 2D projection.

        Parameters
        ----------
        ax : matplotlib.axes.Axes, optional
            Axes to plot on. If None, creates a new figure and axes.
        show_points : bool, optional
            Whether to show the points in the sections.
        show_lines : bool, optional
            Whether to show the lines connecting the points.
        show_domains : bool, optional
            Whether to color sections based on their domain.
        annotate : bool, optional
            Whether to annotate the sections with their index.
        projection : str or tuple, optional
            The projection to use for the plot. Can be 'XY', 'XZ', 'YZ', or a tuple of two axes.
        highlight_sections : list of Section, optional
            Sections to highlight in the plot.
        focus_sections : list of Section, optional
            Sections to focus on in the plot.
        """
        if ax is None:
            fig, ax = plt.subplots(figsize=(10, 10))

        highlight_sections = set(highlight_sections) if highlight_sections else None
        focus_sections = set(focus_sections) if focus_sections else None
        x_attr, y_attr = projection[0].lower(), projection[1].lower()

        section_count = len(self.sections)  # Avoid recalculating

        for sec in self.sections:
            # Skip sections that are not in the focus set (if focus is specified)
            if focus_sections and sec not in focus_sections:
                continue

            xs = [getattr(pt, x_attr) for pt in sec.points]
            ys = [getattr(pt, y_attr) for pt in sec.points]

            # Assign colors based on domains or section index
            color = plt.cm.jet(1 - sec.idx / section_count)
            if show_domains:
                color = get_domain_color(sec.domain)
            if highlight_sections and sec in highlight_sections:
                color = 'red'

            # Plot section points and lines
            if show_points:
                ax.plot(xs, ys, '.', color=color, markersize=7, markeredgecolor='black')
            if show_lines:
                ax.plot(xs, ys, color=color, zorder=0)

            # Annotate section index if needed
            if annotate:
                mean_x, mean_y = np.mean(xs), np.mean(ys)
                ax.annotate(
                    f'{sec.idx}', (mean_x, mean_y), fontsize=8,
                    color='white',
                    bbox=dict(facecolor='black', edgecolor='white', 
                    boxstyle='round,pad=0.3')
                )

        ax.set_xlabel(projection[0])
        ax.set_ylabel(projection[1])
        ax.set_aspect('equal')



    def plot_radii_distribution(self, ax=None, highlight=None, 
    domains=True, show_soma=False):
        """
        Plot the radius distribution of the sections in the tree.

        Parameters
        ----------
        ax : matplotlib.axes.Axes, optional
            Axes to plot on. If None, creates a new figure and axes.
        highlight : list of int, optional
            Indices of sections to highlight in the plot.
        domains : bool, optional
            Whether to color sections based on their domain.
        show_soma : bool, optional
            Whether to show the soma section in the plot.
        """
        if ax is None:
            fig, ax = plt.subplots(figsize=(8, 3))

        for sec in self.sections:
            if not show_soma and sec.parent is None:
                continue
            color = get_domain_color(sec.domain)
            if highlight and sec.idx in highlight:
                ax.plot(
                    [pt.path_distance() for pt in sec.points], 
                    sec.radii, 
                    marker='.', 
                    color='red', 
                    zorder=2
                )
            else:
                ax.plot(
                    [pt.path_distance() for pt in sec.points], 
                    sec.radii, 
                    marker='.', 
                    color=color, 
                    zorder=1
                )
        ax.set_xlabel('Distance from root')
        ax.set_ylabel('Radius')


    # EXPORT METHODS

    def to_swc(self, path_to_file: str):
        """
        Save the SectionTree as an SWC file.

        Parameters
        ----------
        path_to_file : str
            The path to save the SWC file.
        """
        if not self.is_sorted or not self._point_tree.is_sorted:
            raise ValueError('The tree must be sorted before saving.')

        data = {
            'idx': [],
            'type_idx': [],
            'x': [],
            'y': [],
            'z': [],
            'r': [],
            'parent_idx': []
        }

        for sec in self.sections:
            points = sec.points if sec.parent is None or sec.parent.parent is None else sec.points[1:]
            for pt in points:
                data['idx'].append(pt.idx)
                data['type_idx'].append(pt.type_idx)
                data['x'].append(pt.x)
                data['y'].append(pt.y)
                data['z'].append(pt.z)
                data['r'].append(pt.r)
                data['parent_idx'].append(pt.parent_idx)

        df = pd.DataFrame(data)
        df.to_csv(path_to_file, sep=' ', index=False, header=False)

