Source code for pynpoint.readwrite.attr_writing

"""
Modules for writing data as text file.
"""

import os

from typing import Optional

import numpy as np

from typeguard import typechecked

from pynpoint.core.processing import WritingModule


class AttributeWritingModule(WritingModule):
    """
    Module for writing a 1D or 2D array of non-static attributes to a text file.
    """

    __author__ = 'Tomas Stolker'

    @typechecked
    def __init__(self,
                 name_in: str,
                 data_tag: str,
                 attribute: str,
                 file_name: str = 'attributes.dat',
                 output_dir: Optional[str] = None,
                 header: Optional[str] = None) -> None:
        """
        Parameters
        ----------
        name_in : str
            Unique name of the module instance.
        data_tag : str
            Tag of the database entry from which the ``PARANG`` attribute is read.
        attribute : str
            Name of the non-static attribute as given in the central database (e.g., 'INDEX' or
            'STAR_POSITION').
        file_name : str
            Name of the output file.
        output_dir : str, None
            Output directory where the text file will be stored. If no path is specified then the
            Pypeline default output location is used.
        header : str, None
            Header that is written at the top of the text file.

        Returns
        -------
        NoneType
            None
        """

        super().__init__(name_in, output_dir=output_dir)

        self.m_data_port = self.add_input_port(data_tag)

        self.m_file_name = file_name
        self.m_attribute = attribute
        self.m_header = header

    @typechecked
    def run(self) -> None:
        """
        Run method of the module. Writes the non-static attributes (1D or 2D) to a a text file.

        Returns
        -------
        NoneType
            None
        """

        if self.m_header is None:
            self.m_header = ''

        print('Writing attribute data...', end='')

        out_name = os.path.join(self.m_output_location, self.m_file_name)

        if self.m_attribute not in self.m_data_port.get_all_non_static_attributes():
            raise ValueError(f'The \'{self.m_attribute}\' attribute is not present in '
                             f'\'{self.m_data_port.tag}\'.')

        values = self.m_data_port.get_attribute(self.m_attribute)

        np.savetxt(out_name, values, header=self.m_header, comments='# ')

        print(' [DONE]')

        self.m_data_port.close_port()


class ParangWritingModule(WritingModule):
    """
    Module for writing a list of parallactic angles to a text file.
    """

    __author__ = 'Tomas Stolker'

    @typechecked
    def __init__(self,
                 name_in: str,
                 data_tag: str,
                 file_name: str = 'parang.dat',
                 output_dir: Optional[str] = None,
                 header: Optional[str] = None) -> None:
        """
        Parameters
        ----------
        name_in : str
            Unique name of the module instance.
        data_tag : str
            Tag of the database entry from which the ``PARANG`` attribute is read.
        file_name : str
            Name of the output file.
        output_dir : str, None
            Output directory where the text file will be stored. If no path is specified then the
            Pypeline default output location is used.
        header : str, None
            Header that is written at the top of the text file.

        Returns
        -------
        NoneType
            None
        """

        super().__init__(name_in, output_dir=output_dir)

        self.m_data_port = self.add_input_port(data_tag)

        self.m_file_name = file_name
        self.m_header = header

    @typechecked
    def run(self) -> None:
        """
        Run method of the module. Writes the parallactic angles from the ``PARANG`` attribute of
        the specified database tag to a a text file.

        Returns
        -------
        NoneType
            None
        """

        print('Writing parallactic angles...', end='')

        if self.m_header is None:
            self.m_header = ''

        out_name = os.path.join(self.m_output_location, self.m_file_name)

        if 'PARANG' not in self.m_data_port.get_all_non_static_attributes():
            raise ValueError(f'The PARANG attribute is not present in \'{self.m_data_port.tag}\'.')

        parang = self.m_data_port.get_attribute('PARANG')

        np.savetxt(out_name, parang, header=self.m_header, comments='# ')

        print(' [DONE]')

        self.m_data_port.close_port()