# File: thema/thema.py
# Last Update: 07/29/25
# Updated by: JW

"""
Main module for the Thema package.
"""

import os
import networkx as nx
from omegaconf import OmegaConf
from thema.multiverse import Planet, Oort, Galaxy


class Thema:
    """
    A Class to drive all of your Topological Hyperparameter Evaluations and Mapping Needs!

    This is the main entry point for the Thema package, handling the configuration and
    execution of the data processing pipeline from raw data to topology-based analysis.

    Parameters
    ----------
    YAML_PATH : str
        Path to the YAML configuration file containing parameters for the Thema pipeline.

    Attributes
    ----------
    YAML_PATH : str
        The path to the YAML configuration file.
    params : OmegaConf
        The loaded parameters from the YAML file.
    clean_files : list
        List of files generated by the innerSystem processing steps.
    projection_files : list
        List of files generated by the outerSystem processing steps.
    model_files : list
        List of files generated by the galaxy processing steps.

    Examples
    --------
    >>> thema = Thema('params.yaml')
    >>> thema.genesis()  # Run the full pipeline
    >>>
    >>> # Access representative stars after genesis
    >>> selection = thema.galaxy.selection
    >>> print(f"Selected {len(selection)} representative stars")
    >>> for cluster_id, info in selection.items():
    ...     print(f"Cluster {cluster_id}: {info['star']} ({info['cluster_size']} stars)")
    """

    def __init__(self, YAML_PATH):
        """
        Initialize the Thema object with configuration from a YAML file.

        Parameters
        ----------
        YAML_PATH : str
            Path to the YAML configuration file.

        Raises
        ------
        ValueError
            If the YAML file doesn't exist or has incorrect format.
        """
        if os.path.isfile(YAML_PATH) and YAML_PATH.endswith(".yaml"):
            self.YAML_PATH = YAML_PATH
            with open(YAML_PATH, "r") as f:
                self.params = OmegaConf.load(f)
        else:
            raise ValueError("There was an issue with your yaml parameter file")

        self.clean_files = None
        self.projection_files = None
        self.model_files = None
        self.selected_model_files = None

    def genesis(self):
        """
        Run the complete Thema pipeline.

        This method executes all steps of the pipeline in the correct order:
        1. Clean and prepare the innerSystem (data cleaning)
        2. Process the outerSystem (dimension reduction)
        3. Analyze with the galaxy (topology analysis)

        Returns
        -------
        None
            The results are stored in the corresponding directories and
            file lists are updated in the object attributes.

        Examples
        --------
        >>> thema = Thema('params.yaml')
        >>> thema.genesis()
        >>>
        >>> # Representative stars are automatically selected and stored in galaxy.selection
        >>> selected_files = [info['star'] for info in thema.galaxy.selection.values()]
        >>> print(f"Representative files: {selected_files}")
        """
        self.spaghettify_innerSystem()
        self.innerSystem_genesis()

        self.spaghettify_outerSystem()
        self.outerSystem_genesis()

        self.spaghettify_galaxy()
        self.galaxy_genesis()

    def spaghettify(self):
        """
        Clean all output directories from a previous run.

        This method removes all files and folders in the output directory structure,
        providing a clean slate for a new run of the pipeline.

        Returns
        -------
        None
            All output directories are cleaned and removed.

        Examples
        --------
        >>> thema = Thema('params.yaml')
        >>> thema.spaghettify()  # Clean all output directories
        """
        out_dir = os.path.join(self.params.outDir, self.params.runName)
        clean_dir = os.path.join(out_dir, "clean")
        proj_dir = os.path.join(out_dir, "projections")
        model_dir = os.path.join(out_dir, "models")
        self.spaghettify_innerSystem()
        self.spaghettify_outerSystem()
        self.spaghettify_galaxy()

        os.rmdir(clean_dir)
        os.rmdir(proj_dir)
        os.rmdir(model_dir)
        os.rmdir(out_dir)

    def innerSystem_genesis(self):
        """
        Run the innerSystem data cleaning pipeline.

        This method initializes the Planet object to clean and preprocess
        the raw data, generating cleaned data files in the clean directory.

        Returns
        -------
        None
            The results are saved to the clean directory and the clean_files
            attribute is updated with the list of generated files.

        Raises
        ------
        AssertionError
            If the clean directory is not empty.

        Examples
        --------
        >>> thema = Thema('params.yaml')
        >>> thema.spaghettify_innerSystem()  # First clean the directory
        >>> thema.innerSystem_genesis()  # Process the data
        """
        clean_outdir = os.path.join(self.params.outDir, self.params.runName + "/clean")
        if os.path.isdir(clean_outdir):
            assert (
                len(os.listdir(clean_outdir)) == 0
            ), "Your clean data directory is not empty. Please clean it with spaghettify_innerSystem()"
        else:
            os.makedirs(clean_outdir)
        self.planet = Planet(YAML_PATH=self.YAML_PATH)
        self.planet.fit()
        self.clean_files = [clean_outdir + file for file in os.listdir(clean_outdir)]

    def spaghettify_innerSystem(self):
        """
        Clean the innerSystem directory.

        This method removes all files and subdirectories in the clean directory,
        providing a clean slate for a new run of the innerSystem processing.

        Returns
        -------
        None
            The clean directory is emptied of all files.

        Examples
        --------
        >>> thema = Thema('params.yaml')
        >>> thema.spaghettify_innerSystem()  # Clean the inner system directory
        """
        clean_outdir = os.path.join(self.params.outDir, self.params.runName + "/clean")
        if os.path.isdir(clean_outdir):
            for filename in os.listdir(clean_outdir):
                file_path = os.path.join(clean_outdir, filename)
                try:
                    if os.path.isfile(file_path):
                        os.remove(file_path)
                    elif os.path.isdir(file_path):
                        os.rmdir(file_path)
                except Exception as e:
                    print(f"Error while deleting {file_path}: {e}")

    def outerSystem_genesis(self):
        """
        Run the outerSystem dimension reduction pipeline.

        This method initializes the Oort object to perform dimension reduction
        on the cleaned data, generating projection files in the projections directory.

        Returns
        -------
        None
            The results are saved to the projections directory and the projection_files
            attribute is updated with the list of generated files.

        Raises
        ------
        AssertionError
            If the projections directory is not empty.

        Examples
        --------
        >>> thema = Thema('params.yaml')
        >>> thema.spaghettify_outerSystem()  # First clean the directory
        >>> thema.outerSystem_genesis()  # Process the data
        """
        proj_outdir = os.path.join(
            self.params.outDir, self.params.runName + "/projections/"
        )
        if os.path.isdir(proj_outdir):
            assert (
                len(os.listdir(proj_outdir)) == 0
            ), "Your projection directory is not empty. Please clean it with spaghettify_outerSystem()"
        else:
            os.makedirs(proj_outdir)
        self.oort = Oort(YAML_PATH=self.YAML_PATH)
        self.oort.fit()
        self.projection_files = [proj_outdir + file for file in os.listdir(proj_outdir)]

    def spaghettify_outerSystem(self):
        """
        Clean the outerSystem projections directory.

        This method removes all files and subdirectories in the projections directory,
        providing a clean slate for a new run of the outerSystem processing.

        Returns
        -------
        None
            The projections directory is emptied of all files.

        Examples
        --------
        >>> thema = Thema('params.yaml')
        >>> thema.spaghettify_outerSystem()  # Clean the outer system directory
        """
        proj_outdir = os.path.join(
            self.params.outDir, self.params.runName + "/projections/"
        )
        if os.path.isdir(proj_outdir):
            for filename in os.listdir(proj_outdir):
                file_path = os.path.join(proj_outdir, filename)
                try:
                    if os.path.isfile(file_path):
                        os.remove(file_path)
                    elif os.path.isdir(file_path):
                        os.rmdir(file_path)
                except Exception as e:
                    print(f"Error while deleting {file_path}: {e}")

    def galaxy_genesis(self):
        """
        Run the galaxy topology analysis pipeline.

        This method initializes the Galaxy object to perform topological analysis
        on the projected data, generating model files in the models directory.

        Returns
        -------
        None
            The results are saved to the models directory and the model_files
            attribute is updated with the list of generated files.

        Raises
        ------
        AssertionError
            If the models directory is not empty.

        Examples
        --------
        >>> thema = Thema('params.yaml')
        >>> thema.spaghettify_galaxy()  # First clean the directory
        >>> thema.galaxy_genesis()  # Process the data
        >>>
        >>> # Representative stars are stored in galaxy.selection after collapse
        >>> representative_files = [info['star'] for info in thema.galaxy.selection.values()]
        >>> print(f"Found {len(representative_files)} representative stars")
        """
        model_outdir = os.path.join(
            self.params.outDir, self.params.runName + "/models/"
        )
        if os.path.isdir(model_outdir):
            assert (
                len(os.listdir(model_outdir)) == 0
            ), "Your models directory is not empty. Please clean it with spaghettify_galaxy()"
        else:
            os.makedirs(model_outdir)
        self.galaxy = Galaxy(YAML_PATH=self.YAML_PATH)
        self.galaxy.fit()
        self.model_files = [model_outdir + file for file in os.listdir(model_outdir)]
        self.galaxy.collapse()
        self.selected_model_files = [
            str(x["star"]) for x in self.galaxy.selection.values()
        ]

    def spaghettify_galaxy(self):
        """
        Clean the galaxy models directory.

        This method removes all files and subdirectories in the models directory,
        providing a clean slate for a new run of the galaxy processing.

        Returns
        -------
        None
            The models directory is emptied of all files.

        Examples
        --------
        >>> thema = Thema('params.yaml')
        >>> thema.spaghettify_galaxy()  # Clean the galaxy directory
        """
        model_outdir = os.path.join(
            self.params.outDir, self.params.runName + "/models/"
        )
        if os.path.isdir(model_outdir):
            for filename in os.listdir(model_outdir):
                file_path = os.path.join(model_outdir, filename)
                try:
                    if os.path.isfile(file_path):
                        os.remove(file_path)
                    elif os.path.isdir(file_path):
                        os.rmdir(file_path)
                except Exception as e:
                    print(f"Error while deleting {file_path}: {e}")
