Converters and Visualizer

class Converter(ABC):
    """Abstract base class for the converter plugin.

    The abstract base class is used to define the interface for the converter plugins:

    - get_args: Get the arguments from the command line.
    - convert: Convert the input file to the output file.
    - __call__: Call the converter plugin.

    Currently used for:

    - Convertion of the input file.
    - Convertion of the output file.
    """

    @abstractmethod
    def get_args(self) -> Dict[str, Any]:
        """Get the arguments from the command line.

        Returns:
            Dict[str, Any]: Return the input file arguments as a dictionary without
                 additional information beyond the command line arguments.

        Raises:
            ValueError: If the output file format is not supported.
        """

    @staticmethod
    @abstractmethod
    def convert(infile: Path, file_format: str) -> MutableMapping[str, Any]:
        """Convert the input file to the target file format.

        It is an abstract method and must be implemented in the derived class.

        Args:
            infile (Path): Input file as a path object.
            file_format (str): Target file format.

        Returns:
            MutableMapping[str, Any]: Converted file as a dictionary.
        """

    @abstractmethod
    def save(self, data: Any, fname: Path, export_format: str) -> None:
        """Save the data to the target file format.

        Args:
            data (Any): Data to save.
            fname (Path): Filename of the target file.
            export_format (str): Target file format.
        """

    @abstractmethod
    def __call__(self) -> None:
        """Call the converter plugin."""

class FileConverter(Converter):
    """Convert the input and output file to the preferred file format.

    !!! info "Supported file formats"

        Currently supported file formats:

        -[x] JSON
        -[x] YAML (YML)
        -[x] TOML (LOCK for the lock file)
    """

    def get_args(self) -> Dict[str, Any]:
        """Get the arguments from the command line.

        Returns:
            Dict[str, Any]: Return the input file arguments as a dictionary without
                additional information beyond the command line arguments.
        """
        parser = argparse.ArgumentParser(
            description="Converter for 'SpectraFit' input and output files."
        )
        parser.add_argument(
            "infile",
            type=Path,
            help="Filename of the 'SpectraFit' input or output file.",
        )
        parser.add_argument(
            "-f",
            "--file-format",
            help="File format for the conversion.",
            type=str,
            choices=choices,
        )
        parser.add_argument(
            "-e",
            "--export-format",
            help="File format for the export.",
            type=str,
            default="json",
            choices=choices,
        )
        return vars(parser.parse_args())

    @staticmethod
    def convert(infile: Path, file_format: str) -> MutableMapping[str, Any]:
        """Convert the input file to the output file.

        Args:
            infile (Path): The input file as a path object.
            file_format (str): The output file format.

        Raises:
            ValueError: If the input file format is not supported.

        Returns:
            MutableMapping[str, Any] : The converted file as a dictionary.
        """
        if file_format not in choices:
            raise ValueError(f"The input file format '{file_format}' is not supported.")

        return read_input_file(infile)

    def save(self, data: Any, fname: Path, export_format: str) -> None:
        """Save the converted file.

        Raises:
            ValueError: If the input file format is identical with the output format.
            ValueError: If the output file format is not supported.

        Args:
            data (Any): The converted file as a dictionary.
            fname (Path): The input file as a path object.
            export_format (str): The output file format.
        """
        if fname.suffix[1:] == export_format:
            raise ValueError(
                f"The input file suffix '{fname.suffix[1:]}' is similar to the"
                f" output file format '{export_format}'."
                "Please use a different output file suffix."
            )

        if export_format not in choices:
            raise ValueError(
                f"The output file format '{export_format}' is not supported."
            )

        if export_format == "json":
            with open(
                fname.with_suffix(f".{export_format}"), "w", encoding="utf-8"
            ) as f:
                json.dump(data, f, indent=4)
        elif export_format in {"yaml", "yml"}:
            with open(
                fname.with_suffix(f".{export_format}"), "w", encoding="utf-8"
            ) as f:
                yaml.dump(data, f, default_flow_style=False)
        elif export_format in {"toml", "lock"}:
            with open(
                fname.with_suffix(f".{export_format}"),
                "wb+",
            ) as f:
                tomli_w.dump(dict(**data), f)

    def __call__(self) -> None:
        """Run the converter via cmd commands."""
        args = self.get_args()
        self.save(
            data=self.convert(infile=args["infile"], file_format=args["file_format"]),
            fname=args["infile"],
            export_format=args["export_format"],
        )

class DataConverter(Converter):
    """Convert the data files to a CSV file.

    !!! info "Supported file formats"

        Currently supported file formats:

        -[x] ATHENA
        -[x] TXT
        -[ ] more to come

        `DataConverter` class can be also used in the Jupyter notebook.
    """

    def get_args(self) -> Dict[str, Any]:
        """Get the arguments from the command line.

        Returns:
            Dict[str, Any]: Return the input file arguments as a dictionary without
                additional information beyond the command line arguments.
        """
        parser = argparse.ArgumentParser(
            description="Converter for 'SpectraFit' from data files to CSV files."
        )
        parser.add_argument(
            "infile",
            type=Path,
            help="Filename of the data file to convert.",
        )
        parser.add_argument(
            "-f",
            "--file-format",
            help="File format for the conversion.",
            type=str,
            choices=choices,
        )
        parser.add_argument(
            "-e",
            "--export-format",
            help="File format for the export.",
            type=str,
            default="csv",
            choices=choices_export,
        )
        return vars(parser.parse_args())

    @staticmethod
    def convert(infile: Path, file_format: str) -> pd.DataFrame:
        """Convert the input file to the target file format.

        Args:
            infile (Path): Input file as a path object.
            file_format (str): Target file format.

        Raises:
            ValueError: If the file format is not supported.

        Returns:
            pd.DataFrame: The converted data as a pandas DataFrame.
        """
        if file_format.upper() not in choices:
            raise ValueError(f"File format '{file_format}' is not supported.")

        if callable(DataFormats.__dict__[file_format].names):
            names = DataFormats.__dict__[file_format].names(infile)
        else:
            names = DataFormats.__dict__[file_format].names
        DataFormats.__dict__[file_format].names = names

        return pd.read_csv(
            infile, **DataFormats.__dict__[file_format].dict(exclude={"file_suffixes"})
        )

    def save(self, data: Any, fname: Path, export_format: str) -> None:
        """Save the converted data to a CSV file.

        Raises:
            ValueError: If the export format is not supported.

        Args:
            data (Any): The converted data, which is a pandas DataFrame.
            fname (Path): The file name of the data file.
            export_format (str): The file format of the exported file.
        """
        if export_format.lower() not in choices_export:
            raise ValueError(f"Export format '{export_format}' is not supported.")
        data.to_csv(fname.with_suffix(f".{export_format}"), index=False)

    def __call__(self) -> None:
        """Run the converter."""
        args = self.get_args()
        self.save(
            data=self.convert(
                args["infile"],
                args["file_format"],
            ),
            fname=args["infile"],
            export_format=args["export_format"],
        )

@dataclass(frozen=True)
class DataFormats:
    """Data formats."""

    ATHENA = athena_format
    TXT = txt_format

class ExportData:
    """Export the data to a file.

    !!! info "General information"

        The data is exported to a file. The file format is determined by the file
        extension of the output file. The supported file formats are:

        -[x] npy
        -[x] npz
        -[x] pkl
        -[x] pkl.gz

        Classical file formats like `CSV`, `JSON`, `TOML`, etc. are not supported.
        In case of `CSV`, the conversion from unstructured data to a structured
        format is not trivial. In case of `JSON` and `TOML`, the data is not
        the conversion from numpy arrays to lists is very costly. Therefore, the
        data is exported to a pickly file as the preferred format.

    !!! warning "About NumPy"

        The data is exported to a NumPy file can cause some challenge for the
        loading of the data. The data is exported as a dictionary with numpy
        as numpy arrays. The data can be loaded with the following code:

        ```python
        import numpy as np

        data = np.load("data.npy", allow_pickle=True).item()
        ```
    """

    def __init__(self, data: Dict[str, Any], fname: Path, export_format: str) -> None:
        """Export the data to a file.

        Args:
            data (Dict[str, Any]): The data to export.
            fname (Path): The filename of the output file.
            export_format (str): The file format of the output file.
        """
        self.data = data
        self.fname = fname.with_suffix(f".{export_format}")
        self.export_format = export_format

    def __call__(self) -> None:
        """Export the data to a file."""
        if self.export_format in {"npy", "npz"}:
            self.to_numpy()
        elif self.export_format in {"pkl", pkl_gz}:
            self.to_pickle()

    def to_numpy(self) -> None:
        """Export the data to a numpy file."""
        _data: Any = self.data
        if self.export_format.lower() == "npy":
            np.save(self.fname, _data)
        elif self.export_format.lower() == "npz":
            np.savez(self.fname, data=_data)

    def to_pickle(self) -> None:
        """Export the data to a pickle file."""
        if self.export_format.lower() == "pkl":
            with open(self.fname, "wb") as f:
                pickle.dump(self.data, f)
        elif self.export_format.lower() == pkl_gz:
            with gzip.open(self.fname, "wb") as f:
                pickle.dump(self.data, f)

    @staticmethod
    def numpy2list(data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Convert the arrays of list dictionaries to a list of dictionaries with list.

        Args:
            data (List[Dict[str, Any]]): The data to convert.

        Returns:
            List[Dict[str, Any]]: The converted data.
        """
        return [
            {k: v.tolist() for k, v in d.items() if isinstance(v, np.ndarray)}
            for d in data
        ]

class PklConverter(Converter):
    """Convert pkl data to a CSV files.

    !!! info "General information"

        The pkl data is converted to a CSV file. The CSV file is saved in the same
        directory as the input file. The name of the CSV file is the same as the
        input file with the suffix `.csv` and prefixed with the name of the
        'major' keys in the pkl file. Furthermore, a graph of the data is optionally
        saved as a PDF file to have a visual representation of the data structure.

    !!! info "Supported file formats"

        Currently supported file formats:

        -[x] pkl
        -[x] pkl.gz
        -[x] ...

    """

    def get_args(self) -> Dict[str, Any]:
        """Get the arguments from the command line.

        Returns:
            Dict[str, Any]: Return the input file arguments as a dictionary without
                additional information beyond the command line arguments.
        """
        parser = argparse.ArgumentParser(
            description="Converter for 'SpectraFit' from pkl files to CSV files."
        )
        parser.add_argument(
            "infile",
            type=Path,
            help="Filename of the pkl file to convert.",
        )
        parser.add_argument(
            "-f",
            "--file-format",
            help="File format for the optional encoding of the pickle file."
            " Default is 'latin1'.",
            type=str,
            default="latin1",
            choices=choices_fformat,
        )
        parser.add_argument(
            "-e",
            "--export-format",
            help="File format for export of the output file. Default is 'pkl'.",
            type=str,
            default="pkl",
            choices=choices_export,
        )
        return vars(parser.parse_args())

    @staticmethod
    def convert(infile: Path, file_format: str) -> Dict[str, Any]:
        """Convert the input file to the output file.

        Args:
            infile (Path): The input file of the as a path object.
            file_format (str): The output file format.

        Returns:
            Dict[str, Any]: The data as a dictionary, which can be a nested dictionary
        """

        def _convert(
            data_values: Dict[str, Any], _key: Optional[List[str]] = None
        ) -> List[Dict[str, Any]]:
            """Convert the data to a list of dictionaries.

            The new key is the old key plus all the subkeys. The new value is the
            value of the subkey if the value is an instance of an array.

            For avoiding `pylint` errors, the `_key` argument is set to `None` by
            default and is set to an empty list if it is `None`. This is done to
            avoid the `pylint` error `dangerous-default-value`. The `_key` argument
            is used to keep track of the keys of the nested dictionary. Furthermore,
            the `_key` argument is used to create the new key for the new dictionary.

            Finally, the new dictionary is appended to the list of dictionaries.

            Args:
                data_values (Dict[str, Any]): The data as a dictionary.

            Returns:
                List[Dict[str, Any]]: The data as a list of dictionaries.
            """
            data_list = []
            if _key is None:
                _key = []
            for key, value in data_values.items():
                if isinstance(value, dict):
                    _key.append(str(key))
                    data_list.extend(_convert(value, _key))
                    _key.pop()
                elif isinstance(value, np.ndarray):
                    data_list.append({"_".join(_key + [key]): value})
            return data_list

        data_dict = {}
        for key, value in pkl2any(infile, file_format).items():
            if isinstance(value, dict):
                data_dict[key] = _convert(value)
        return data_dict

    def save(self, data: Any, fname: Path, export_format: str) -> None:
        """Save the converted pickle data to a file.

        Args:
            data (Any): The converted nested dictionary of the pkl data.
            fname (Path): The filename of the output file.
            export_format (str): The file format of the output file.

        Raises:
            ValueError: If the export format is not supported.
        """
        if export_format.lower() not in choices_export:
            raise ValueError(f"Unsupported file format '{export_format}'.")

        fname = pure_fname(fname)

        for key, value in data.items():
            _fname = Path(f"{fname}_{key}").with_suffix(f".{export_format}")
            ExportData(data=value, fname=_fname, export_format=export_format)()

    def __call__(self) -> None:
        """Run the converter."""
        args = self.get_args()
        data = self.convert(args["infile"], args["file_format"])
        self.save(data, args["infile"], args["export_format"])

class PklVisualizer(Converter):
    """Visualize the pkl data as a graph."""

    def get_args(self) -> Dict[str, Any]:
        """Get the arguments from the command line.

        Returns:
            Dict[str, Any]: Return the input file arguments as a dictionary without
                additional information beyond the command line arguments.
        """
        parser = argparse.ArgumentParser(
            description="Converter for 'SpectraFit' from pkl files to a graph."
        )
        parser.add_argument(
            "infile",
            type=Path,
            help="Filename of the pkl file to convert to graph.",
        )
        parser.add_argument(
            "-f",
            "--file-format",
            help="File format for the optional encoding of the pickle file."
            " Default is 'latin1'.",
            type=str,
            default="latin1",
            choices=choices_fformat,
        )
        parser.add_argument(
            "-e",
            "--export-format",
            help="File extension for the graph export.",
            type=str,
            default="pdf",
            choices=choices_export,
        )

        return vars(parser.parse_args())

    @staticmethod
    def convert(infile: Path, file_format: str) -> Dict[str, Any]:
        """Convert the input file to the output file.

        Args:
            infile (Path): The input file of the as a path object.
            file_format (str): The encoding of the pickle file.

        Raises:
            ValueError: If the data is not a dictionary.

        Returns:
            Dict[str, Any]: The data as a dictionary, which can be a nested dictionary.
        """
        data = PklVisualizer().get_type(pkl2any(infile, encoding=file_format))
        if not isinstance(data, dict):
            raise ValueError(f"Data is not a dictionary: {data}")
        graph = PklVisualizer().create_graph(fname=infile, data_dict=data)

        pos = nx.kamada_kawai_layout(graph, scale=2)
        nx.draw_networkx_nodes(
            graph, pos, node_size=100, node_color="lightblue", alpha=0.8
        )
        nx.draw_networkx_edges(graph, pos, width=0.5, edge_color="grey", alpha=0.5)
        nx.draw_networkx_labels(graph, pos, font_size=10, font_family="sans-serif")
        plt.axis("off")
        return data

    def save(self, data: Any, fname: Path, export_format: str) -> None:
        """Save the graph to a file and the data and their types to a json file.

        Args:
            data (Any): The data to save, which can be a nested dictionary.
            fname (Path): The filename of the file to save.
            export_format (str): The file format to save the graph to.

        Raises:
            ValueError: If the export format is not supported.
        """
        if export_format.lower() not in choices_export:
            raise ValueError(f"Export format '{export_format}' is not supported.")

        plt.savefig(
            pure_fname(fname).with_suffix(f".{export_format}"),
            format=export_format,
        )

        with open(
            pure_fname(fname).with_suffix(".json"), "w+", encoding="utf-8"
        ) as outfile:
            json.dump(data, outfile, indent=4)

    def get_type(self, value: Any) -> Union[Dict[str, Any], str]:
        """Get the type of the value.

        Args:
            value (Any): The value to get the type from.

        Returns:
            Union[Dict[str, Any], str]: The type of the value.
        """
        if isinstance(value, dict):
            return {key: self.get_type(value) for key, value in value.items()}
        if isinstance(value, np.ndarray):
            return f"{type(value)} of shape {value.shape}"
        return str(type(value))

    def add_nodes(self, graph: nx.DiGraph, data_dict: Dict[str, Any]) -> None:
        """Add nodes to the graph.

        Args:
            graph (nx.DiGraph): The graph to add nodes to.
            data_dict (Dict[str, Any]): The data dictionary to get the nodes from.
        """
        for key, value in data_dict.items():
            graph.add_node(key)
            if isinstance(value, dict):
                for item in value:
                    graph.add_edge(key, item)
                    graph.add_node(item)
                self.add_nodes(graph=graph, data_dict=value)
            elif "of shape" in str(value):
                value = value.split("of shape")
                graph.add_node(value[0])
                graph.add_edge(key, value[0])

                graph.add_node(value[-1])
                graph.add_edge(value[0], value[-1])

            else:
                graph.add_node(value)
                graph.add_edge(key, value)

    def create_graph(self, fname: Path, data_dict: Dict[str, Any]) -> nx.DiGraph:
        """Create the graph.

        Args:
            fname (Path): The filename of the file to create the graph from.
            data_dict (Dict[str, Any]): The data dictionary to create the graph from.

        Returns:
            nx.DiGraph: The graph created from the data dictionary.
        """
        graph = nx.DiGraph()
        graph.add_node(str(fname.name))
        for key in data_dict:
            graph.add_edge(str(fname.name), key)

        self.add_nodes(graph=graph, data_dict=data_dict)
        return graph

    def __call__(self) -> None:
        """Create the graph and save it as a PDF file."""
        args = self.get_args()
        self.save(
            data=self.convert(args["infile"], args["file_format"]),
            fname=args["infile"],
            export_format=args["export_format"],
        )
        plt.show()

class RIXSConverter(Converter):
    """Transform the raw pkl data into a JSON, TOML, or numpy file for RIXS."""

    def get_args(self) -> Dict[str, Any]:
        """Get the arguments from the command line.

        Returns:
            Dict[str, Any]: Return the input file arguments as a dictionary without
                additional information beyond the command line arguments.
        """
        parser = argparse.ArgumentParser(
            description="Converter for 'SpectraFit' from pkl files to a JSON, TOML, "
            "or numpy file for RIXS-Visualizer."
        )
        parser.add_argument(
            "infile",
            type=Path,
            help="Filename of the pkl file to convert to JSON, TOML, or numpy.",
        )
        parser.add_argument(
            "-f",
            "--file-format",
            help="File format for the optional encoding of the pickle file."
            " Default is 'latin1'.",
            type=str,
            default="latin1",
            choices=choices_fformat,
        )
        parser.add_argument(
            "-e",
            "--export-format",
            help="File extension for the export.",
            type=str,
            default="json",
            choices=choices_export,
        )
        parser.add_argument(
            "-ie",
            "--incident_energy",
            help="Name of the incident energy",
            type=str,
        )
        parser.add_argument(
            "-ee",
            "--emission_energy",
            help="Name of the emitted energy",
            type=str,
        )
        parser.add_argument(
            "-rm",
            "--rixs_map",
            help="Name of the RIXS map",
        )
        parser.add_argument(
            "-m",
            "--mode",
            help="Mode of the RIXS map post-processing, e.g. 'sum' or 'max'."
            "Default is 'sum'.",
            type=str,
            default="sum",
            choices=choices_mode,
        )
        return vars(parser.parse_args())

    @staticmethod
    def convert(infile: Path, file_format: str) -> MutableMapping[str, Any]:
        """Convert the pkl file to a dictionary.

        Args:
            infile (Path): The input file.
            file_format (str): The file format for the optional encoding of the pickle
                file.

        Returns:
            MutableMapping[str, Any]: The data dictionary from the pkl file.
        """
        data_dict = {}
        for _dict in pkl2any(infile, file_format):
            data_dict.update(_dict)
        return data_dict

    def create_rixs(
        self,
        data: MutableMapping[str, Any],
        incident_energy: str,
        emission_energy: str,
        rixs_map: str,
        mode: str,
    ) -> RIXSModelAPI:
        """Create the RIXS map from the pkl file.

        Args:
            data (MutableMapping[str, Any]): The data dictionary from the pkl file.
            incident_energy (str): The name of the incident energy.
            emission_energy (str): The name of the emitted energy.
            rixs_map (str): The name of the RIXS map.
            mode (str): The mode of the RIXS map post-processing, e.g. 'sum' or 'max'.

        Raises:
            ValueError: If the mode is not in the choices.
            KeyError: If the incident energy, emission energy, or RIXS map is not in
                the data.

        Returns:
            RIXSModelAPI: The RIXS map as a RIXSModelAPI pydantic object.
        """
        if mode not in choices_mode:
            raise ValueError(f"Mode '{mode}' not in {choices_mode}.")
        if incident_energy not in data:
            self.raise_error(incident_energy, data)
        if emission_energy not in data:
            self.raise_error(incident_energy, data)
        if rixs_map not in data:
            self.raise_error(incident_energy, data)

        if mode == "sum":
            rixs_val = np.sum(data[rixs_map], axis=0)
        elif mode == "mean":
            rixs_val = np.mean(data[rixs_map], axis=0)
        return RIXSModelAPI(
            incident_energy=data[incident_energy],
            emission_energy=data[emission_energy],
            rixs_map=rixs_val,
        )

    @staticmethod
    def raise_error(wrong_key: str, data: Any) -> None:
        """Raise an error if the key is not in the data.

        Args:
            wrong_key (str): The key which is not in the data.
            data (Any): The data dictionary from the pkl file.

        Raises:
            KeyError: If the key is not in the data.

        """
        raise KeyError(
            f"Key '{wrong_key}' not in data. Aailable keys are: {list(data.keys())}."
        )

    def save(self, data: Any, fname: Path, export_format: str) -> None:
        """Save the data to a file.

        Args:
            data (Any): The data to save.
            fname (Path): The filename.
            export_format (str): The file extension for the export.

        Raises:
            ValueError: If the export format is not in the choices.
        """
        if export_format not in choices_export:
            raise ValueError(
                f"Export format '{export_format}' not in {choices_export}."
            )

        if export_format == "json":
            with open(
                pure_fname(fname).with_suffix(f".{export_format}"),
                "w",
                encoding="utf-8",
            ) as f:
                json.dump(self.numpydict2listdict(data), f, indent=4)
        elif export_format in {"toml", "lock"}:
            with open(
                pure_fname(fname).with_suffix(f".{export_format}"),
                "wb",
            ) as f:
                tomli_w.dump(self.numpydict2listdict(data), f, multiline_strings=False)
        elif export_format == "npy":
            np.save(pure_fname(fname).with_suffix(f".{export_format}"), data)
        elif export_format == "npz":
            np.savez(pure_fname(fname).with_suffix(f".{export_format}"), **data)

    @staticmethod
    def numpydict2listdict(data: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
        """Convert a dictionary with numpy arrays to a dictionary with lists.

        Args:
            data (MutableMapping[str, Any]): The data dictionary.

        Returns:
            MutableMapping[str, Any]: The data dictionary with lists.
        """
        return {k: v.tolist() for k, v in data.items()}

    def __call__(self) -> None:
        """Run the converter."""
        args = self.get_args()
        self.save(
            data=self.create_rixs(
                data=self.convert(args["infile"], args["file_format"]),
                incident_energy=args["incident_energy"],
                emission_energy=args["emission_energy"],
                rixs_map=args["rixs_map"],
                mode=args["mode"],
            ).dict(),
            fname=args["infile"],
            export_format=args["export_format"],
        )