Source code for susan.data.Reference

###########################################################################
# This file is part of the Substack Analysis (SUSAN) framework.
# Copyright (c) 2018-2021 Ricardo Miguel Sanchez Loayza.
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
# 
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
###########################################################################

import susan.utils.txt_parser as _prsr
from susan.utils import is_extension as _is_ext
from susan.utils import force_extension as _force_ext



class Reference:
    """Container for one or more reference maps used in STA alignment.

    Each reference entry holds paths to four MRC files: the full map,
    a soft mask, and the two independent half-maps.  The half-maps are
    used for FSC-based resolution estimation and as N2N training pairs.

    File format: plain-text ``.refstxt``.

    .. rubric:: Attributes

    Attributes
    ----------
    ref : list of str
        Paths to the full reference maps.
    msk : list of str
        Paths to the soft masks (same box size as the references).
    h1  : list of str
        Paths to half-map 1 for each reference.
    h2  : list of str
        Paths to half-map 2 for each reference.
    """

    ref = []
    msk = []
    h1  = []
    h2  = []

    def __init__(self, filename=None, n_refs=0):
        """Load from file or allocate an empty container.

        Parameters
        ----------
        filename : str, optional
            Path to a ``.refstxt`` file to load.
        n_refs : int, optional
            Number of empty reference slots to allocate (used when
            filename is None).
        """
        self.ref.clear()
        self.msk.clear()
        self.h1.clear()
        self.h2.clear()
        
        if isinstance(filename,str):
            fp = open(filename,"rb")
            num_refs = int(_prsr.parse_args(fp)['num_ref'])
            for i in range(num_refs):
                args = _prsr.parse_args(fp)
                self.ref.append(args['map' ])
                self.msk.append(args['mask'])
                self.h1.append( args['h1'  ])
                self.h2.append( args['h2'  ])
        elif( n_refs>0 ):
            for i in range(n_refs):
                self.ref.append('')
                self.msk.append('')
                self.h1.append('')
                self.h2.append('')
    
    def get_n_refs(self):
        """Return the number of references stored."""
        return len(self.ref)
    
    n_refs = property(get_n_refs)
    
    @staticmethod
    def _check_filename(filename):
        if not _is_ext(filename,'refstxt'):
            raise ValueError( 'Wrong file extension, do you mean ' + _force_ext(filename,'refstxt') + '?')
    
    #def __repr__(self):
    #    return "Reference"
    


    @staticmethod
    def load(filename):
        """Load a Reference from a ``.refstxt`` file.

        Parameters
        ----------
        filename : str
            Path to a ``.refstxt`` file.

        Returns
        -------
        Reference
        """
        Reference._check_filename(filename)
        result = Reference(filename=filename)
        return result

    


    def save(self, filename):
        """Save to a ``.refstxt`` file.

        Parameters
        ----------
        filename : str
            Output path; must have a ``.refstxt`` extension.
        """
        Reference._check_filename(filename)
        
        if len(self.ref) is not len(self.msk):
            raise NameError('Reference entries do not match Mask entries')
            
        if len(self.ref) is not len(self.h1):
            raise NameError('Reference entries do not match Half1 entries')

        if len(self.ref) is not len(self.h2):
            raise NameError('Reference entries do not match Half2 entries')
        
        fp=open(filename,'w')
        _prsr.write(fp,'num_ref',str(self.n_refs))
        for i in range(self.n_refs):
            fp.write('## Reference '+str(i+1)+'\n')
            _prsr.write(fp,'map' ,self.ref[i])
            _prsr.write(fp,'mask',self.msk[i])
            _prsr.write(fp,'h1'  ,self.h1[i] )
            _prsr.write(fp,'h2'  ,self.h2[i] )
        fp.close()