from collections import defaultdict
from multiprocessing import Pool, cpu_count

import numpy as np
import pandas as pd
import seaborn as sns
from matplotlib import pyplot as plt
from rdkit.Chem import Draw, DataStructs, BondType, AllChem as Chem
from tqdm import tqdm

BOND_DICT = {
    BondType.SINGLE: "-",
    BondType.DOUBLE: "=",
    BondType.TRIPLE: "#",
    BondType.AROMATIC: ":",
}


class FingerprintManager:
    """
    Fingerprint class
    """

    def __init__(self):
        self.subs_per_on_bits = None
        self.spo_filtered = None
        self.frags = None
        self.fp = None
        self.bit_info_name = None
        self.fp_name = None
        self.n_bits = None
        self.radius = None
        self.use_features = None
        self.input_df = None
        self.mol_col = None

    @staticmethod
    def fetch_fp_from_mol(
        in_mol: Chem.Mol, radius: int = 2, n_bits: int = 2048, use_features: bool = False, use_chirality: bool = False
    ) -> tuple[np.ndarray, dict[str, list[int]]]:
        """
        Generate Morgan fingerprint from a RDKit molecule object.

        Parameters
        ----------
        in_mol : Chem.Mol
            RDKit molecule object
        radius : int, optional
            Radius of the circular fingerprint (default is 2).
        n_bits : int, optional
            Number of bits in the fingerprint (default is 2048).
        use_features : bool, optional
            Use Feature Morgan fingerprint or not (default is False).
        use_chirality : bool, optional
            Use chirality or not (default is False).

        Returns
        -------
        tuple[np.ndarray, dict[str, list[int]]]
            Returning a two-element tuple containing the fingerprint as a numpy array as the first element and
            bit info as a dictionary as the second element.

        """
        bit_info = {}
        morgan_fp = Chem.GetMorganFingerprintAsBitVect(
            in_mol,
            radius,
            nBits=n_bits,
            useFeatures=use_features,
            bitInfo=bit_info,
            useChirality=use_chirality,
        )
        morgan_fp_as_array = np.empty(n_bits, dtype=np.uint8)
        DataStructs.ConvertToNumpyArray(morgan_fp, morgan_fp_as_array)
        return morgan_fp_as_array, bit_info

    @staticmethod
    def fetch_fp_from_smiles(
        in_smiles: str,
        radius: int = 2,
        n_bits: int = 2048,
        use_features: bool = False,
        use_chirality: bool = False,
    ) -> tuple[np.ndarray, dict[str, list[int]]]:
        """
        Generate Morgan fingerprint from a SMILES string.

        Parameters
        ----------
        in_smiles : str
            SMILES string
        radius : int, optional
            Radius of the circular fingerprint (default is 2).
        n_bits : int, optional
            Number of bits in the fingerprint (default is 2048).
        use_features : bool, optional
            Use Feature Morgan fingerprint or not (default is False).
        use_chirality : bool, optional
            Use chirality or not (default is False).

        Returns
        -------
        morgan_fp_as_array : tuple[np.ndarray, dict[str, list[int]]]
            Returning a two-element tuple containing the fingerprint as a numpy array as the first element and
            bit info as a dictionary as the second element.

        Raises
        ------
        ValueError
            If the SMILES string is invalid.

        """
        # Catch errors based on invalid SMILES
        try:
            mol = Chem.MolFromSmiles(Chem.CanonSmiles(in_smiles))
            return FingerprintManager.fetch_fp_from_mol(mol, radius, n_bits, use_features, use_chirality)
        except:
            raise ValueError(f"Invalid SMILES present: {in_smiles}")

    def fetch_fp_from_df(
        self,
        input_df: pd.DataFrame,
        smiles_col: str = "SMILES",
        radius: int = 2,
        n_bits: int = 2048,
        use_features: bool = False,
        use_chirality: bool = False,
        mol_col: bool = None,
        canonicalize: bool = True,
    ) -> None:
        """
        Generate either Morgan Fingerprints or Feature Morgan Fingerprints. Bit info is always included.

        Parameters
        ----------
        input_df : pd.DataFrame
            Pandas DataFrame containing SMILES strings.
        smiles_col : str, optional
            Name of the column containing SMILES strings (default is "SMILES").
        radius : int, optional
            Radius of the circular fingerprint (default is 2).
        n_bits : int, optional
            Number of bits in the fingerprint (default is 2048).
        use_features : bool, optional
            Use Feature Morgan fingerprint or not (default is False).
        use_chirality : bool, optional
            Use chirality or not (default is False).
        mol_col : str, optional
            Name of the column containing RDKit molecules (default is None).
        canonicalize : bool, optional
            If True, molecules are canonicalized (default is True).

        Raises
        ------
        ValueError
            If either of the SMILES or Molecules column is not present in the input DataFrame.
        """
        self.input_df = input_df
        self.use_features = use_features
        self.radius = radius
        self.n_bits = n_bits
        self.mol_col = mol_col

        # Check whether the SMILES or mol column is present in the input DataFrame
        if smiles_col not in self.input_df.columns and (
            self.mol_col is None or self.mol_col not in self.input_df.columns
        ):
            raise ValueError("Target column not found in the input DataFrame")

        # Check whether the Molecules column is present in the input DataFrame, if no, then make it
        if self.mol_col is None:
            self.mol_col = "Molecules"

            # generate canonical RDKit molecules in the input DataFrame. These molecules will be used many times.
            self.input_df[self.mol_col] = self.input_df[smiles_col].apply(
                lambda x: Chem.MolFromSmiles(Chem.CanonSmiles(x))
            )
        elif canonicalize:  # molecules are present, will canonicalize them
            self.input_df[self.mol_col] = self.input_df[self.mol_col].apply(
                lambda x: Chem.MolFromSmiles(Chem.MolToSmiles(x))
            )

        # Generate Morgan fingerprints and bit info
        self.fp_name = f"Morgan_FP_{self.radius}_{self.n_bits}"
        self.bit_info_name = f"Bit_Info_{self.radius}_{self.n_bits}"
        self.input_df[self.fp_name], self.input_df[self.bit_info_name] = zip(
            *self.input_df[self.mol_col].apply(
                lambda x: self.fetch_fp_from_mol(x, self.radius, self.n_bits, self.use_features, use_chirality)
            )
        )

        # Saving the fingerprints as an array
        self.fp = np.stack(self.input_df[self.fp_name])

    @staticmethod
    def _gen_frag_for_single_bit(
        subs_per_on_bits: list[tuple[Chem.Mol, int, int]], include_connectivity_invariants: bool = True
    ) -> tuple[list[str], list[tuple[Chem.Mol, int, int]]]:
        """
        Function generates all fragments for a single bit including canonical SMARTS patterns of each fragment.

        Parameters
        ----------
        subs_per_on_bits : list[tuple[Chem.Mol, int, int]]
            A list of tuples containing the molecule, atom index and radius each. The list contains all molecules for
            which one particular bit is set.
        include_connectivity_invariants : bool, optional
            If True, connectivity invariants are included in the fragments (default is True).

        Returns
        -------
        tuple[list[str], list[tuple[Chem.Mol, int, int]]]
            Returning a two-element tuple containing a list of unique canonical SMARTS patterns as first element and a
            list of tuples containing a molecule, atom index and radius each as second element.
        """
        # Check if global "ici" variable is set, if so, use it for the invariant parameter
        if "_ici" in globals():
            include_connectivity_invariants = _ici

        seen = []
        spo_filtered = []
        for mol, atom_index, radius in subs_per_on_bits:
            env = Chem.FindAtomEnvironmentOfRadiusN(mol, radius, atom_index, useHs=True)
            atoms = {atom_index}
            for bidx in env:
                atoms.add(mol.GetBondWithIdx(bidx).GetBeginAtomIdx())
                atoms.add(mol.GetBondWithIdx(bidx).GetEndAtomIdx())

            outer_atoms = set()
            if include_connectivity_invariants:
                env = Chem.FindAtomEnvironmentOfRadiusN(mol, radius + 1, atom_index, useHs=True)
                atoms_rpp = {atom_index}
                for bidx in env:
                    atoms_rpp.add(mol.GetBondWithIdx(bidx).GetBeginAtomIdx())
                    atoms_rpp.add(mol.GetBondWithIdx(bidx).GetEndAtomIdx())
                outer_atoms = atoms_rpp - atoms
                atoms = atoms_rpp

            atom_symbols = []
            for atom in mol.GetAtoms():
                if atom.GetIdx() in outer_atoms:
                    atom_symbols.append("*")
                    continue
                atom_symbol = Chem.MolFragmentToSmiles(mol, atomsToUse=[atom.GetIdx()], allHsExplicit=True)
                atom_symbol = f"{atom_symbol[:-1]}{';R' if atom.IsInRing() else ''}]"
                atom_symbols.append(atom_symbol)

            bond_symbols = [
                "~"
                if bond.GetBeginAtomIdx() in outer_atoms or bond.GetEndAtomIdx() in outer_atoms
                else BOND_DICT[bond.GetBondType()]
                for bond in mol.GetBonds()
            ]

            smi = Chem.MolFragmentToSmiles(
                mol,
                atomsToUse=list(atoms),
                bondsToUse=env,
                atomSymbols=atom_symbols,
                bondSymbols=bond_symbols,
                allBondsExplicit=True,
                rootedAtAtom=atom_index,
                canonical=True,
            )
            if smi not in seen:
                # Check if substructure encoded by "smi" is already present in "seen" but with different atom order
                matches = set()
                for known_smi in seen:
                    matches.add(tuple(sorted(mol.GetSubstructMatch(Chem.MolFromSmarts(known_smi)))))
                if not tuple(sorted(mol.GetSubstructMatch(Chem.MolFromSmarts(smi)))) in matches:
                    spo_filtered.append((mol, atom_index, radius))
                    seen.append(smi)
        return seen, spo_filtered

    @staticmethod
    def _init_pool(ici: bool) -> None:
        """
        Helper function to be able to use the Pool method imap() instead of starmap() because of better performance. It
        takes a boolean specifying if connectivity invariants should be included or not and sets it to the global
        variable `_ici`.

        Parameters
        ----------
        ici : bool
            Boolean specifying if connectivity invariants should be included or not.
        """
        global _ici
        _ici = ici

    def generate_frags(
        self, verbose: bool = True, n_jobs: int = 1, include_connectivity_invariants: bool = True
    ) -> None:
        """
        Generates fragments for all molecules and all nonzero bits in the dataset saved to this instance. The dataset
        must be saved to this instance using the `fetch_fp_from_df()` method.

        Parameters
        ----------
        verbose : bool, optional
            If True, status messages are printed (default is True).
        n_jobs : int, optional
            Number of cores to use for parallelization (default is 1).
        include_connectivity_invariants : bool, optional
            If True, connectivity invariants are included in the fragments (default is True).

        Raises
        ------
        ValueError
            If no input data is present, i.e. `fetch_fp_from_df()` was not called already.
        """
        # Check if there is input data
        if self.input_df is None:
            raise ValueError("Input data is not present. Please use fetch_fp_from_df() method to add input data.")

        # Check n_jobs parameter
        if n_jobs < 1:
            n_jobs = 1
        elif n_jobs > cpu_count():
            n_jobs = cpu_count()

        if verbose:
            print("Generating fragments. This step may take a while.")

        # Gather information for all the bit sets over the entire dataset
        self.subs_per_on_bits = defaultdict(list)

        for i, (_, row) in enumerate(self.input_df.iterrows()):
            single_mol = row[self.mol_col]
            single_mol_bit_info = row[f"Bit_Info_{self.radius}_{self.n_bits}"]
            for bit in single_mol_bit_info:
                for atom_index, radius in single_mol_bit_info[bit]:
                    self.subs_per_on_bits[bit].append((single_mol, atom_index, radius))

        # Print the bits which are set
        if verbose:
            print(f"{len(self.subs_per_on_bits)} / {self.n_bits} bits are set in the dataset.")

        # Filtering duplicates per bit and extract fragment SMARTS
        if n_jobs == 1:
            self.spo_filtered = {}
            self.frags = {}
            for bit in tqdm(self.subs_per_on_bits):
                self.frags[bit], self.spo_filtered[bit] = self._gen_frag_for_single_bit(
                    self.subs_per_on_bits[bit], include_connectivity_invariants
                )
        else:
            if verbose:
                print(f"Parallelize using {n_jobs} cores.")

            with Pool(n_jobs, initializer=self._init_pool, initargs=(include_connectivity_invariants,)) as p:
                frags, spo_filtered = zip(*p.imap(self._gen_frag_for_single_bit, self.subs_per_on_bits.values()))

            self.frags = dict(zip(self.subs_per_on_bits.keys(), frags))
            self.spo_filtered = dict(zip(self.subs_per_on_bits.keys(), spo_filtered))

    def make_substruct_per_bit_plot(self, dpi: int = 100, save_fig: bool = False, fig_name: str = None) -> None:
        """
        Generate substructures per bit plot. This plot is useful to understand the distribution of substructures encoded
        in each bit.

        Parameters
        ----------
        dpi : int, optional
            Dots per inch (default is 100).
        save_fig : bool, optional
            If True, the figure is saved (default is False).
        fig_name : str, optional
            Name of the plot (default is None).
        """
        # Check for frags and spo_filtered
        if self.frags is None or self.spo_filtered is None:
            print("WARNING: Fragments were not generated, yet. Try to generate now with default settings...")
            self.generate_frags()

        plt.figure(figsize=(2, 4), dpi=dpi)
        fig = sns.boxplot([len(self.frags[bit]) for bit in self.frags])
        fig.set_ylabel("Substructures per bit")
        fig.set_xlabel(f'{"F" if self.use_features else "E"}CFP{self.radius * 2} {self.n_bits} bits')

        if save_fig:
            if fig_name is None:
                fig_name = f"Substructures_per_bit_plot.png"
            plt.savefig(fig_name, dpi=dpi)
        else:
            plt.show()

    def substruct_counter_per_bit(self, query_bit: int, verbose: bool = False) -> pd.DataFrame:
        """
        Counts the frequency of the substructures present in a query bit. Total occurrences will calculate all the
        occurrences of a substructure in the dataset, including multiple occurrences in a molecule. Unique occurrences
        will calculate the number of molecules in which a substructure is present.

        Parameters
        ----------
        query_bit : int
            Bit for which the substructure frequency is to be calculated.
        verbose : bool, optional
            If True, status messages are printed (default is False).

        Returns
        -------
        pd.DataFrame
            Pandas DataFrame where substructures (as SMARTS) are present with their total and unique occurrences.
        """
        # Check for frags and spo_filtered
        if self.frags is None or self.spo_filtered is None:
            print("WARNING: Fragments were not generated, yet. Try to generate now with default settings...")
            self.generate_frags()

        # Initialize an empty list to store unique occurrences.
        substruct_unique_occurrences = []

        # Initialize an empty list to store total occurrences.
        substruct_total_occurrences = []

        # Print wait-time warning
        if verbose:
            print(f"Counting substructure frequencies of the query bit: Bit {query_bit}. This step may take a while.")

        # Run a loop of all the substructures over the data to prepare their count
        for substructure in self.frags[query_bit]:
            frequency_substructure_per_mol = 0
            frequency_substructure_per_mol_unique = 0
            substructure_as_mol = Chem.MolFromSmarts(substructure)

            # Loop over molecules in the dataset
            for single_mol in self.input_df[self.mol_col]:
                matches = single_mol.GetSubstructMatches(substructure_as_mol)
                frequency_substructure_per_mol += len(matches)

                # To calculate unique occurrences
                if len(matches) > 0:
                    frequency_substructure_per_mol_unique += 1

            substruct_total_occurrences.append(frequency_substructure_per_mol)
            substruct_unique_occurrences.append(frequency_substructure_per_mol_unique)

        # Creating the substructure dataframe
        substruct_df = pd.DataFrame(
            {
                "Serial": range(1, len(self.frags[query_bit]) + 1),
                "Substructures (as SMARTS)": self.frags[query_bit],
                "Total occurrences": substruct_total_occurrences,
                "Unique occurrences": substruct_unique_occurrences,
            }
        )
        return substruct_df

    def get_substruct_images_per_bit(self, query_bit: int, smarts_as_legends: bool = False) -> str:
        """
        Generates an SVG image of substructures present in a query bit.

        Parameters
        ----------
        query_bit : int
            Bit for which the substructure images are to be generated.
        smarts_as_legends : bool, optional
            If True, substructure SMARTS are used as legends (default is False).

        Returns
        -------
        str
            SVG image as a string of substructures present in a query bit. This image should be saved separately.

        """
        # Check for frags and spo_filtered
        if self.frags is None or self.spo_filtered is None:
            print("WARNING: Fragments were not generated, yet. Try to generate now with default settings...")
            self.generate_frags()

        # Fix Draw Option issue:
        draw_options = Draw.rdMolDraw2D.MolDrawOptions()
        draw_options.prepareMolsBeforeDrawing = False

        # legends
        if smarts_as_legends:
            legends = self.frags[query_bit]
        else:
            legends = [str(i + 1) for i in range(len(self.frags[query_bit]))]

        # Get substructure images
        substruct_images = Draw.DrawMorganEnvs(
            self.spo_filtered[query_bit],
            legends=legends,
            subImgSize=(300, 300),
            molsPerRow=4,
            drawOptions=draw_options,
        )
        return substruct_images

src/xfp/core_xfp.py
czodrowskilab/x-fp
import importlib.resources as pkg_resources
from datetime import datetime

import pandas as pd
from fpdf import FPDF, Align

from . import resources


class BitAnalysis(FPDF):
    def header(self) -> None:
        """
        Generate header of the report.
        """

        # For the cover page, add the X-FP logo and the CzodrowskiLab logo
        if self.page_no() == 1:
            # X-FP Logo
            self.image(
                pkg_resources.files(resources) / "X-FP_logo.png",
                w=60,
                x=Align.L,
                y=10,
            )

            # CzodrowskiLab Logo
            self.image(
                pkg_resources.files(resources) / "CzodrowskiLab_Logo.png",
                w=30,
                x=Align.R,
                y=5,
            )

        else:
            self.title = "X-FP"

            # Setting header's font:
            self.set_font("helvetica", "B", 12)

            # Calculating width of title
            width = self.get_string_width(self.title) + 6

            # Setting cursor position:
            self.set_x((210 - width) / 2)

            # Setting thickness of the frame (1 mm)
            self.set_line_width(1)
            # Printing title:
            self.cell(
                width,
                9,
                self.title,
                new_x="LMARGIN",
                new_y="NEXT",
                align="C",
                fill=False,
            )

            self.ln(5)

            # printing subtitle
            self.set_font("helvetica", "I", 10)
            self.set_x((210 - width) / 2)
            self.cell(
                width,
                h=5,
                txt="Bit Analysis",
                new_x="LMARGIN",
                new_y="NEXT",
                align="C",
                fill=False,
            )

        # Performing a line break:
        self.ln(10)

    def footer(self) -> None:
        """
        Generate footer of the report.
        """
        # Setting position at 1.5 cm from bottom:
        self.set_y(-15)
        # Setting font: helvetica italic 8
        self.set_font("helvetica", "I", 8)
        # Setting text color to gray:
        self.set_text_color(128)
        # Adding the user provided report title on the left side of the footer
        self.cell(0, 10, self.report_title, align="L")

        # Setting position at 1.5 cm from bottom:
        self.set_y(-15)
        # Setting font: helvetica italic 8
        self.set_font("helvetica", "I", 8)
        # Setting text color to gray:
        self.set_text_color(128)
        # Printing page number
        self.cell(0, 10, f"Page {self.page_no()} of {{nb}}", align="C")

        # time stamp
        if self.time_stamp:
            # Setting position at 1.5 cm from bottom:
            self.set_y(-15)
            # Setting font: helvetica italic 8
            self.set_font("helvetica", "I", 8)
            # Setting text color to gray:
            self.set_text_color(128)
            # Printing time stamp
            self.cell(0, 10, f"Report generated on {self.time_stamp}", align="R")

    def intro_page(
        self, overall_feature_importance_image: str, img_text: str, report_title: str = "", time_stamp: bool = True
    ) -> None:
        """
        Generate an introduction page of the report where the overall feature importance plot is displayed.

        Parameters
        ----------
        overall_feature_importance_image : str
            Image of the overall feature importance plot.
        img_text : str
            Text to be displayed above the overall feature importance image.
        report_title : str, optional
            Title of the report provided by the user (default is "").
        time_stamp : bool, optional
            Add time stamp to the report footer or not (default is True). The time is local and is formatted as
            '%Y-%m-%d %H:%M:%S'.
        """

        # Truncate the title length to 50 characters
        if len(report_title) > 50:
            report_title = f"{report_title[:50]}..."

        # Creating a class variable to store the report title
        self.report_title = report_title

        # Creating a class variable to store the time stamp
        if time_stamp:
            self.time_stamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        else:
            self.time_stamp = None

        # Add page
        self.add_page()
        self.set_y(70)  # in case report title is missing

        # writing the title of the report if provided
        if report_title != "":
            self.set_font("helvetica", "U", 20)
            self.set_y(40)
            self.multi_cell(w=0, txt=self.report_title, align="C")
            self.ln(20)

        self.image(
            overall_feature_importance_image,
            w=150,
            x=Align.C,
        )
        self.ln(5)

        # adding the title of the plot
        self.set_font("helvetica", "", 10)
        self.multi_cell(0, 5, img_text, align="C")
        self.ln(20)

    def chapter_title(self, in_bit: int) -> None:
        """
        Generate chapter title of the report.

        Parameters
        ----------
        in_bit : int
            Bit of choice for which the analysis is to be performed
        """

        # Setting font: helvetica 12
        self.set_font("helvetica", "", 12)
        # Setting background color
        self.set_fill_color(200, 220, 255)
        # Printing chapter name:
        self.cell(
            0,
            6,
            # f"Chapter {num} : {label}",
            f"Bit {in_bit}",
            new_x="LMARGIN",
            new_y="NEXT",
            align="L",
            fill=True,
        )
        # Performing a line break:
        self.ln(4)

    def print_chapter(
        self,
        in_bit: int,
        mol_importance_image: str,
        mol_image: str,
        in_df: pd.DataFrame,
        mol_importance_text: str = None,
        n_compounds: int = None,
    ) -> None:
        """
        Generate chapter of the report.

        Parameters
        ----------
        in_bit : int
            Bit of choice for which the analysis is to be performed.
        mol_importance_image : str
            Image of the feature importance plot for the query bit.
        mol_image : str
            Image of the substructure plot for the query bit.
        in_df : pd.DataFrame
            Pandas DataFrame containing the substructure frequency table for the query bit.
        mol_importance_text : str, optional
            Text to be displayed above the feature importance image (default is None).
        n_compounds : int, optional
            Total number of compounds in the input dataframe (default is None).
        """

        # Setting font: helvetica 8
        self.set_font("helvetica", "", 8)

        self.add_page()
        self.chapter_title(in_bit)

        # Feature importance image, e.g., SHAP Plot

        if mol_importance_text is None:
            mol_importance_text = f"SHAP Summary Plot for the Bit {in_bit}:"

        self.multi_cell(0, 5, mol_importance_text)
        self.ln()
        self.image(mol_importance_image, w=100, x=20)
        self.ln()

        # Substructure images
        sub_struct_heading_text = f"Substructures present:"
        self.multi_cell(0, 5, sub_struct_heading_text)
        self.ln()
        self.image(mol_image, w=150, x=20)
        self.ln()

        # Substructure table

        if n_compounds is None:
            sub_struct_table_text = "Substructures frequency table:"
        else:
            sub_struct_table_text = f"Substructures frequency table (total compounds: {n_compounds}):"
        self.multi_cell(0, 5, sub_struct_table_text)
        self.ln()

        in_df = in_df.copy()
        in_df = in_df.applymap(str)  # Convert all data inside dataframe into string type

        columns = [list(in_df)]  # Get list of dataframe columns
        rows = in_df.values.tolist()  # Get list of dataframe rows
        data = columns + rows  # Combine columns and rows in one list

        self.set_font("helvetica", size=8)
        with self.table(
            borders_layout="HORIZONTAL_LINES",
            # cell_fill_color=200,  # grey
            cell_fill_mode="ROWS",
            line_height=self.font_size * 2.5,
            text_align="CENTER",
            width=180,
            col_widths=(20, 100, 30, 30),
        ) as table:
            for data_row in data:
                row = table.row()
                for datum in data_row:
                    row.cell(datum)
        self.ln()

    def add_end_note(self, feature_importance_note: str = None) -> None:
        """
        Generate end note of the report.

        Parameters
        ----------
        feature_importance_note : str, optional
            Path to the text file containing the information about feature importance (default is None).
        """

        # Add page
        self.add_page()

        # setting font of title
        self.set_font("helvetica", "B", 12)

        self.cell(
            0,
            6,
            f"Notes about substructures rendering: ",
            new_x="LMARGIN",
            new_y="NEXT",
            align="L",
            # fill=True,
        )

        # setting font of text
        self.set_font("helvetica", "", 8)

        # Reading text file of substructure rendering:
        substructure_rendering_file = pkg_resources.files(resources) / "rendering_info.txt"
        with substructure_rendering_file.open() as fh:
            substructure_rendering_text = fh.read()  # .decode("latin-1")

        self.multi_cell(0, 5, substructure_rendering_text)
        self.ln()

        # Adding note about feature importance
        if feature_importance_note:
            self.set_font("helvetica", "B", 12)

            self.cell(
                0,
                6,
                f"Notes about feature importance: ",
                new_x="LMARGIN",
                new_y="NEXT",
                align="L",
                # fill=True,
            )

            # setting font of text
            self.set_font("helvetica", "", 8)

            # Reading text file of feature importance:
            if feature_importance_note == "shap_tree_explainer":
                feature_importance_file = pkg_resources.files(resources) / "shap_tree_explainer_info.txt"

            else:
                feature_importance_file = feature_importance_note

            with feature_importance_file.open() as fh:
                feature_importance_text = fh.read()

            self.multi_cell(0, 5, feature_importance_text)
            self.ln()

src/xfp/report_maker.py
czodrowskilab/x-fp
import datetime
import os
from typing import Optional

import matplotlib.pyplot as plt
import numpy as np
import shap

from .core_xfp import FingerprintManager
from .report_maker import BitAnalysis


class XFPTreeExplainer(FingerprintManager):
    def __init__(self, in_tree_explainer: shap.TreeExplainer) -> None:
        """
        Initialize the XFPTreeExplainer class.

        Parameters
        ----------
        in_tree_explainer : shap.TreeExplainer
            Tree explainer object from SHAP to use.

        Raises
        ------
        ValueError
            If the input tree explainer is not a valid SHAP Tree Explainer.
        """
        super().__init__()
        self.shap_values = None
        self.tree_explainer = in_tree_explainer

        # Check whether the tree explainer is valid or not
        if not isinstance(self.tree_explainer, shap.TreeExplainer):
            raise ValueError("Invalid SHAP Tree Explainer")

    def generate_shap_values(self, return_shap_values: bool = False) -> Optional[np.ndarray]:
        """
        Generate SHAP values of the input fingerprint using the Tree Explainer.

        Parameters
        ----------
        return_shap_values : bool, optional
            Return the SHAP values or not (default is False).

        Returns
        -------
        np.ndarray or None
            SHAP values of the input fingerprint in `self.fp` if `return_shap_values` is True, None otherwise.
        """
        self.shap_values = self.tree_explainer.shap_values(self.fp)

        if return_shap_values:
            return self.shap_values

    def _folder_settings(self) -> None:
        """
        Asks for input for folder (with path) to store the images and the report.
        Validates the path and creates the folder if it doesn't exist.
        """

        # Ask for input for folder (with path) to store the images and the report
        self.folder_name = input(
            f"Enter the folder name (with path) to store the images and the report. The default is 'XFP_Tree_Explainer_(today's date)' "
        )

        # Validate the path and create the folder if it doesn't exist
        if self.folder_name == "":
            # Add date and a folder name
            self.folder_name = f"XFP_Tree_Explainer_{datetime.date.today()}"

        if not os.path.exists(self.folder_name):
            os.makedirs(self.folder_name)

    def generate_shap_summary_plot(self, max_display: int = 10, save_fig: str = "", dpi: int = 300) -> None:
        """
        Generate a SHAP summary plot.

        Parameters
        ----------
        max_display : int, optional
            Maximum number of features to plot (default is 10).
        save_fig : str, optional
            If provided, save the figure instead of showing it (default is "").
        dpi : int, optional
            Dots per inch (DPI) for the figure (default is 300).
        """

        # Feature names
        feature_names = [f"Bit {i}" for i in range(self.fp.shape[1])]

        # Check if save_fig is provided, if yes, save the figure instead of showing it.
        if save_fig != "":
            shap.summary_plot(
                self.shap_values,
                self.fp,
                max_display=max_display,
                feature_names=feature_names,
                show=False,
            )
            plt.savefig(save_fig, dpi=dpi, bbox_inches="tight")
            plt.close()
            return

        shap.summary_plot(self.shap_values, self.fp, max_display=max_display, feature_names=feature_names)

    def generate_single_shap_summary_plot(self, query_bit: int, return_fig_path: bool = True) -> Optional[str]:
        """
        Saves a single SHAP summary plot of the SHAP values of the query bit.

        Parameters
        ----------
        query_bit : int
            Bit for which the SHAP Summary Plot is to be generated.
        return_fig_path : bool, optional
            Return the path of the figure or not (default is True).

        Returns
        -------
        str or None
            Path of the figure if `return_fig_path` is True, None otherwise.
        """

        # Making the summary plot of only the query bit.
        shap.summary_plot(
            self.shap_values[:, query_bit : query_bit + 1],
            self.fp[:, query_bit : query_bit + 1],
            feature_names=[f"Bit {query_bit}"],
            show=False,
        )

        # Save the plot
        shap_summary_plot_path = f"{self.folder_name}/Bit_{query_bit}_SHAP_Summary_Plot.png"
        plt.savefig(shap_summary_plot_path, dpi=1000, bbox_inches="tight")

        plt.clf()
        plt.close()

        if return_fig_path:
            return shap_summary_plot_path

    def generate_bit_analysis_report(
        self,
        bit_list: list[int],
        report_title: str = "",
        time_stamp: bool = True,
        file_name: str = "X-FP_Bit_Analysis_Report",
    ) -> None:
        """
        Generate a PDF report for the analysis of the query bits.

        Parameters:
        -----------
        bit_list : list[int]
            List of bits for which the analysis is to be performed.
        report_title : str, optional
            Title of the report provided by the user (default is "").
        time_stamp : bool, optional
            Add time stamp to the report footer or not (default is True). The time is local and is formatted as
            '%Y-%m-%d %H:%M:%S'.
        file_name : str, optional
            Name of the report file (default is "X-FP_Bit_Analysis_Report").
        """

        # folder settings
        self._folder_settings()

        # Start the report
        bit_analysis_report = BitAnalysis()

        # Generate the overall feature importance plot
        self.generate_shap_summary_plot(save_fig=f"{self.folder_name}/Overall_Feature_Importance_Plot.png")
        overall_shap_summary_plot_path = f"{self.folder_name}/Overall_Feature_Importance_Plot.png"

        # Print the title page
        bit_analysis_report.intro_page(
            overall_feature_importance_image=overall_shap_summary_plot_path,
            img_text="SHAP Summary Plot for the Top 10 Morgan Fingerprint Bits",
            report_title=report_title,
            time_stamp=time_stamp,
        )

        # Looping over all the query bits
        for query_bit in bit_list:
            # Generate the SHAP summary plot for the query bit
            shap_summary_plot_path = self.generate_single_shap_summary_plot(query_bit, return_fig_path=True)

            # Generating the count of substructures for the query bit, and retrieving the dataframe
            substructure_count_df = self.substruct_counter_per_bit(query_bit)

            # Generating the substructure images for the query bit, and saving that image in the folder
            substructure_images = self.get_substruct_images_per_bit(query_bit)
            substructure_images_path = f"{self.folder_name}/Bit_{query_bit}_Substructure_Images.svg"
            with open(substructure_images_path, "w") as svg_file:
                svg_file.write(substructure_images)

            # Generating the bit analysis report
            bit_analysis_report.print_chapter(
                in_bit=query_bit,
                mol_importance_image=shap_summary_plot_path,
                mol_image=substructure_images_path,
                in_df=substructure_count_df,
                n_compounds=len(self.input_df),
            )

        # Add end note to the report
        bit_analysis_report.add_end_note(feature_importance_note="shap_tree_explainer")

        # Save the report
        bit_analysis_report.output(f"{self.folder_name}/{file_name}.pdf")

src/xfp/xfp_tree_explainer.py
czodrowskilab/x-fp
import datetime
import importlib.resources as pkg_resources
import os
from typing import Optional

import matplotlib.pyplot as plt
import pandas as pd
import xgboost as xgb
from xgboost import XGBClassifier

from . import resources
from .core_xfp import FingerprintManager
from .report_maker import BitAnalysis


class XGBFP(FingerprintManager):
    def __init__(self, in_xgboost_model: xgb.XGBClassifier) -> None:
        """
        Initialize the XGBFP class.

        Parameters
        ----------
        in_xgboost_model : xgboost.XGBClassifier
            XGBoost Classifier model to use.

        Raises
        ------
        ValueError
            If the input model is not a valid XGBoost Classifier model.
        """
        super().__init__()
        self.xgboost_model = in_xgboost_model

        # Check whether the model is a valid XGBoost Classifier model or not
        if not isinstance(self.xgboost_model, xgb.XGBClassifier):
            raise ValueError("Invalid object. Please input valid XGBoost Classifier model")

    def _folder_settings(self) -> None:
        """
        Asks for input for folder (with path) to store the images and the report.
        Validates the path and creates the folder if it doesn't exist.
        """

        # Ask for input for folder (with path) to store the images and the report
        self.folder_name = input(
            f"Enter the folder name (with path) to store the images and the report. The default is 'XFP_XGBoost_Classifier_(today's date)' "
        )

        # Validate the path and create the folder if it doesn't exist
        if self.folder_name == "":
            # Add date and a folder name
            self.folder_name = f"XFP_XGBoost_Classifier_{datetime.date.today()}"

        if not os.path.exists(self.folder_name):
            os.makedirs(self.folder_name)

    @staticmethod
    def _get_feature_importance(
        in_model: XGBClassifier,
        importance_type: str = "weight",
        n_bits: int = 2048,
    ) -> list[float]:
        """
        Get feature importance of the XGBoost Classifier model.

        Parameters
        ----------
        in_model : xgboost.XGBClassifier
            XGBoost Classifier model
        importance_type : str, optional
            Type of importance to be calculated (default is "weight"). The options are "weight", "gain", "cover",
            "total_gain", and "total_cover".
        n_bits : int, optional
            Number of bits in the fingerprint (default is 2048).

        Returns
        -------
        list[float]
            List of feature importance values for all the bits in the fingerprint.

        Notes
        -----
        This function is necessary because there isn't a straight-forward way to get feature importance from a XGBoost
        model for all (query) features. XGBoost doesn't store zero-importance features, plus there are different feature
        importance types to be considered.
        """

        # Assign names to the bits in the model object
        bit_names = [f"Bit {bit}" for bit in range(n_bits)]
        in_model.get_booster().feature_names = bit_names

        # Get feature importance
        importance_score = []
        for bit_name in bit_names:
            try:
                # XGBoost doesn't store zero-importance features
                importance_score.append(in_model.get_booster().get_score(importance_type=importance_type)[bit_name])
            except KeyError:
                importance_score.append(0)

        return importance_score

    def make_feature_importance_df(
        self,
        importance_type: str = None,
        calculate_all_importances: bool = False,
    ) -> None:
        """
        Make a DataFrame containing queried XGBoost feature importance.

        Parameters
        ----------
        importance_type : str, optional
            Type of importance to be calculated (default is None). The options are "weight", "gain", "cover",
            "total_gain", and "total_cover".
        calculate_all_importances : bool, optional
            Additionally calculate all the importances (default is False).

        Raises
        ------
        ValueError
            If the user doesn't provide either importance_type or calculate_all_importances or if the provided
            importance_type is invalid.
        """
        self.importance_df = pd.DataFrame()

        # validate if the user entered either importance type or calculate_all_importances
        if importance_type is None and not calculate_all_importances:
            raise ValueError("Please either enter an importance type or calculate all the importances.")

        # validate if the user entered correct importance type
        if importance_type not in ["weight", "gain", "cover", "total_gain", "total_cover"]:
            raise ValueError("Invalid importance type. Please enter a valid importance type.")

        self.importance_df["Bits"] = [f"Bit {bit}" for bit in range(self.n_bits)]
        if calculate_all_importances:
            importance_types = ["weight", "gain", "cover", "total_gain", "total_cover"]
            print(f"Calculating all importances: {importance_types}. Please wait.")
            for importance_type in importance_types:
                self.importance_df[importance_type] = self._get_feature_importance(
                    self.xgboost_model, importance_type=importance_type, n_bits=self.n_bits
                )
            print("Done.")

        else:
            self.importance_df[importance_type] = self._get_feature_importance(
                self.xgboost_model, importance_type=importance_type, n_bits=self.n_bits
            )

        if importance_type:
            self.primary_importance_type = importance_type

    def generate_feaure_importance_plot(
        self, importance_type: str = None, max_importance: int = 10, save_fig: str = ""
    ) -> None:
        """
        Generate the feature importance plot.

        Parameters
        ----------
        importance_type : str, optional
            Type of importance to be used for plotting (default is None).
            If the primary_importance_type is already set, then that will be used for plotting.
        max_importance : int, optional
            Maximum number of bits to be plotted (default is 10).
        save_fig : str, optional
            Path to save the figure (default is "", i.e. the image is not saved).

        Raises
        ------
        ValueError
            If the user hasn't run the `.make_feature_importance_df()` method first, or if the user doesn't provide an
            importance type or has not set the primary_importance_type, or if the user provides an invalid importance
            type.

        Notes
        -----
        Displays the plot if save_fig is not provided.
        """

        # validate if self.importance_df exists
        if self.importance_df is None:
            raise ValueError("Please run the make_feature_importance_df() method first.")

        # Ensuring that an importance type is selected for plotting
        if self.primary_importance_type is None:
            if importance_type is None:
                raise ValueError("Please enter an importance type.")
            elif importance_type not in self.importance_df.columns:
                raise ValueError("Invalid importance type. Please enter a valid importance type.")
            else:
                self.primary_importance_type = importance_type

        # Sort the dataframe by the primary importance type
        self.importance_df.sort_values(by=self.primary_importance_type, ascending=False, inplace=True)

        # Plotting the feature importance
        fig, ax = plt.subplots()
        ax.barh(
            y=self.importance_df["Bits"][:max_importance],
            width=self.importance_df[self.primary_importance_type][:max_importance],
            color="blue",
        )
        ax.bar_label(ax.containers[0], fmt="%.1f")  # labels on each bar
        ax.margins(x=0.1)  # padding the y-margin so that the numbers fit in properly

        # inverting to avoid ascending order
        plt.gca().invert_yaxis()

        # Hide the left, right and top spines
        ax.spines[["left", "right", "top"]].set_visible(False)

        # Adding the title
        # plt.title(f"Top {max_importance} Bits by XGBoost's {self.primary_importance_type.capitalize()} Feature Importance")

        # Adding labels
        plt.ylabel("")
        plt.xlabel(f"f-score (feature importance by {self.primary_importance_type.capitalize()})")

        # check if save_fig is provided
        if save_fig != "":
            plt.savefig(save_fig, bbox_inches="tight", dpi=300)
            plt.clf()
            plt.close()
            return
        else:
            plt.show()

    def single_feature_importance_plot(self, query_bit: int, return_fig_path: bool = True) -> Optional[str]:
        """
        Saves a feature importance plot for a single bit.

        Parameters
        ----------
        query_bit : int
            Bit for which the feature importance plot is to be generated.
        return_fig_path : bool, optional
            Whether to return the path to the figure or not (default is True).

        Returns
        -------
        str or None
            Path to the figure if `return_fig_path` is True, None otherwise.

        Raises
        ------
        ValueError
            If the user hasn't run the `.make_feature_importance_df()` method first.

        Notes:
        -----
        The plot is saved in the folder with the name:
        Bit_{query_bit}_Feature_Importance_{self.primary_importance_type}_Plot.png
        """
        # validate if self.importance_df exists
        if self.importance_df is None:
            raise ValueError("Please run the make_feature_importance_df() method first.")

        # making the feature importance plot of only the query bit
        fig, ax = plt.subplots(figsize=[8, 1])
        ax.barh(
            y=self.importance_df["Bits"][query_bit],
            width=self.importance_df[self.primary_importance_type][query_bit],
            color="blue",
        )
        ax.bar_label(ax.containers[0], fmt="%.1f")  # labels on each bar
        ax.margins(x=0.1)  # padding the y-margin so that the numbers fit in properly

        # hide the left, right and top spines
        ax.spines[["left", "right", "top"]].set_visible(False)

        # adding the title
        # plt.title(f"Top {max_importance} Bits by XGBoost's {self.primary_importance_type.capitalize()} Feature Importance")

        # adding labels
        plt.ylabel("")
        plt.xlabel(f"f-score (feature importance by {self.primary_importance_type.capitalize()})")

        # save the plot
        feature_importance_plot_path = (
            f"{self.folder_name}/Bit_{query_bit}_Feature_Importance_{self.primary_importance_type}_Plot.png"
        )
        plt.savefig(feature_importance_plot_path, dpi=1000, bbox_inches="tight")

        plt.clf()
        plt.close()

        if return_fig_path:
            return feature_importance_plot_path

    def generate_bit_analysis_report(
        self,
        bit_list: list[int],
        report_title: str = "",
        time_stamp: bool = True,
        file_name: str = "X-FP_Bit_Analysis_Report",
    ) -> None:
        """
        Generates a bit analysis PDF report for the query bits.

        Parameters
        ----------
        bit_list : list[int]
            List of bits for which the bit analysis report is to be generated.
        report_title : str, optional
            Title of the report (default is "").
        time_stamp : bool, optional
            Whether to add a time stamp to the report or not (default is True).
        file_name : str, optional
            Name of the report file (default is "X-FP_Bit_Analysis_Report").

        Raises
        ------
        ValueError
            If the user hasn't run the `.make_feature_importance_df()` method first.
        """

        # folder settings
        self._folder_settings()

        # validate if self.importance_df exists
        if self.importance_df is None:
            raise ValueError("Please run the make_feature_importance_df() method first.")

        # start the report
        bit_analysis_report = BitAnalysis()

        # generate the overall feature importance plot
        self.generate_feaure_importance_plot(
            save_fig=f"{self.folder_name}/Overall_Feature_Importance_{self.primary_importance_type}_Plot.png"
        )
        overall_feature_importance_image = (
            f"{self.folder_name}/Overall_Feature_Importance_{self.primary_importance_type}_Plot.png"
        )

        # print the title page
        bit_analysis_report.intro_page(
            overall_feature_importance_image=overall_feature_importance_image,
            img_text=f"XGBoost's {self.primary_importance_type.capitalize()} Feature Importance Plot for the Top 10 Morgan Fingerprint Bits",
            report_title=report_title,
            time_stamp=time_stamp,
        )

        # looping through the bits
        for query_bit in bit_list:
            # generate the feature importance plot for the query bit, saving plot, and retrieving the path
            single_summary_plot_path = self.single_feature_importance_plot(query_bit=query_bit, return_fig_path=True)

            # Generating the count of substructures for the query bit, and retrieving the dataframe
            substructure_count_df = self.substruct_counter_per_bit(query_bit)

            # Generating the substructure images for the query bit, and saving that image in the folder
            substructure_images = self.get_substruct_images_per_bit(query_bit)
            substructure_images_path = f"{self.folder_name}/Bit_{query_bit}_Substructure_Images.svg"
            with open(substructure_images_path, "w") as svg_file:
                svg_file.write(substructure_images)

            # Generating the bit analysis report
            bit_analysis_report.print_chapter(
                in_bit=query_bit,
                mol_importance_image=single_summary_plot_path,
                mol_image=substructure_images_path,
                in_df=substructure_count_df,
                mol_importance_text=f"XGBoost '{self.primary_importance_type.capitalize()}' feature importance plot for the Bit {query_bit}",
                n_compounds=len(self.input_df),
            )

        # Add end note to the report
        path_to_xgb_feature_importance_note = pkg_resources.files(resources) / "xgb_feature_importance.txt"
        bit_analysis_report.add_end_note(feature_importance_note=path_to_xgb_feature_importance_note)

        # Save the report
        bit_analysis_report.output(f"{self.folder_name}/{file_name}.pdf")

src/xfp/xfp_xgb.py
czodrowskilab/x-fp