# SPDX-FileCopyrightText: 2025 Poirazi Lab <dendrotweaks@dendrites.gr>
# SPDX-License-Identifier: MPL-2.0

import pandas as pd

SWC_IDS_TO_DOMAINS = {
    0: 'undefined',
    1: 'soma',
    2: 'axon',
    3: 'dend',
    31: 'basal',
    4: 'apic',
    41: 'trunk',
    42: 'tuft',
    43: 'oblique',
    6: 'neurite',
    7: 'glia',
    8: 'reduced',
}

DOMAINS_TO_COLORS = {
    'undefined': '#7F7F7F',
    'soma': '#E69F00',
    'axon': '#F0E442',
    'dend': '#019E73',
    'basal': '#31A354',
    'apic': '#0072B2',
    'trunk': '#56B4E9',
    'tuft': '#A55194',
    'oblique': '#8C564B',
    'neurite': '#D55E00',
    'glia': '#D62728',
    'reduced': '#E377C2',
}

DEFAULT_IDS = list(SWC_IDS_TO_DOMAINS.keys())

class SWCReader():
    """
    Reads an SWC file and returns a DataFrame.
    """

    def __init__(self):
        pass


    def read_file(self, path_to_swc_file: str) -> pd.DataFrame:
        """
        Read the SWC file and return a DataFrame enriched with domain & color metadata.
        """
        with open(path_to_swc_file, 'r') as f:
            raw_lines = f.readlines()

        # 1. Clean up lines
        lines = self._clean_lines(raw_lines)

        # 2. Read SWC table
        df = self._read_swc_table(path_to_swc_file)
        unique_type_ids = df['Type'].unique()

        # 3. Extract data from header (if available)
        domain_map, color_map = self._build_mappings(lines, unique_type_ids)

        # 4. Add domain and color columns
        df.insert(df.columns.get_loc("Type") + 1, "Domain", df["Type"].map(domain_map))
        df.insert(df.columns.get_loc("Domain") + 1, "Color", df["Domain"].map(color_map))

        return df
        

    @staticmethod
    def rename_domain(df, type_idx: int, new_domain_name: str, new_color: str):
        """
        Rename a domain in the DataFrame by changing the domain name and color for a given type index.
        """
        df.loc[df['Type'] == type_idx, 'Domain'] = new_domain_name
        df.loc[df['Type'] == type_idx, 'Color'] = new_color


    @staticmethod
    def replace_domain(df, old_type_idx: int, new_type_idx: int):
        """
        Replace a domain in the DataFrame by changing the type index, domain name, and color.

        Parameters
        ----------
        df : pd.DataFrame
            The DataFrame containing the SWC data (generated by read_file).
        old_type_idx : int
            The type index of the domain to be replaced.
        new_type_idx : int
            The type index of the new domain.
        """
        new_domain_name = df.loc[df['Type'] == new_type_idx, 'Domain'].iloc[0]
        new_color = df.loc[df['Type'] == new_type_idx, 'Color'].iloc[0]
        df.loc[df['Type'] == old_type_idx, 'Type'] = new_type_idx
        df.loc[df['Type'] == new_type_idx, 'Domain'] = new_domain_name
        df.loc[df['Type'] == new_type_idx, 'Color'] = new_color

    @staticmethod
    def remove_domain(df, type_idx: int):
        """
        Remove a domain from the DataFrame by dropping all rows with the given type index.
        """
        df.drop(df[df['Type'] == type_idx].index, inplace=True)


    @staticmethod
    def _clean_lines(lines):
        """Strip whitespace, collapse multiple spaces, drop empty lines."""
        return [' '.join(line.split()) for line in lines if line.strip()]


    @staticmethod
    def _read_swc_header(lines):
        """
        Extract domain and color mappings from SWC header.
        """
        domain_map = {}
        color_map = {}

        domain_line = next((l for l in lines if l.startswith("# DOMAIN_NAMES")), None)
        color_line  = next((l for l in lines if l.startswith("# DOMAIN_COLORS")), None)

        if domain_line:
            items = domain_line.replace("# DOMAIN_NAMES", "").strip().split()
            for item in items:
                t, dom = item.split(":", 1)
                domain_map[int(t)] = dom

        if color_line:
            items = color_line.replace("# DOMAIN_COLORS", "").strip().split()
            for item in items:
                t, col = item.split(":", 1)
                domain_name = domain_map.get(int(t))
                color_map[domain_name] = col

        return domain_map, color_map


    def _build_mappings(self, lines, unique_type_ids):
        """
        Build domain and color mappings from SWC header or use defaults.
        """
        domain_map = SWC_IDS_TO_DOMAINS.copy()
        color_map = DOMAINS_TO_COLORS.copy()
        
        # Extract and apply header mappings
        header_domains, header_colors = self._read_swc_header(lines)
        domain_map.update(header_domains)
        color_map.update(header_colors)
        
        # Handle missing type IDs (reduced and custom domains)
        for t in unique_type_ids:
            if t in domain_map:
                continue
            domain_map[t] = f'custom_{t}'
            color_map[f'custom_{t}'] = '#17BECF'
        
        return domain_map, color_map


    @staticmethod
    def _read_swc_table(path):
        """Read the numeric part of an SWC file into a DataFrame."""
        df = pd.read_csv(
            path,
            sep=' ',
            header=None,
            comment='#',
            names=['Index', 'Type', 'X', 'Y', 'Z', 'R', 'Parent'],
            index_col=False,
            dtype={
                'Index': int,
                'Type': int,
                'X': float,
                'Y': float,
                'Z': float,
                'R': float,
                'Parent': int
            }
        )

        # Fix zero-radii case:
        if (df['R'] == 0).all():
            df['R'] = 1.0

        if df['Index'].duplicated().any():
            raise ValueError("The SWC file contains duplicate node ids.")

        return df


    @staticmethod
    def plot_raw_data(df, ax, **kwargs):
        """
        Plot the raw data from the SWC file using the DataFrame.

        Parameters
        ----------
        df : pd.DataFrame
            The DataFrame containing the SWC data (generated by read_file).
        ax : matplotlib.pyplot.Axes
            The axes to plot on.
        """
        for t in df['Type'].unique():
            mask = df['Type'] == t
            color = df[mask]['Color'].iloc[0]
            if pd.isna(color):
                color = 'black'
            ax.scatter(xs=df[mask]['X'], ys=df[mask]['Y'], zs=df[mask]['Z'], 
                    c=color, label=f'Type {t}', **kwargs)
        ax.legend()