#!/usr/bin/python3

# orthologer - Pipeline for deducing orthology from fasta files.
# Copyright (C) 2005-2023 E. Zdobnov lab: F. Tegenfeldt
# <fredrik.tegenfeldt@unige.ch> based on code by E. Kriventseva
#
# This file is part of orthologer.
#
# orthologer is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# orthologer is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with orthologer.  If not, see <http://www.gnu.org/licenses/>.
'''
Check fasta file.
Pipeline usage: PREPROC
'''
from __future__ import print_function

#
# This code checks for a given fasta file if the sequences are proteins or nucletides.
#
import sys
import shutil
import argparse
import fasta_tools as FT
import orthologer_tools as OT
import tempfile as Tmp

def getFastaStats( fasta, dosh, perlbl, verbose ):
    # loop over all fasta
    done=False
    errmsg=None
    ntot=0
    nskip=0
    nerror=0
    minlen=None
    maxlen=None
    sumlen=0.0
    summlen=0.0
    minid=None
    maxid=None
    stype=None
    labels=None
    masked=None

    while not done:
        ntot=ntot+1

        # get next fasta
        (eof,label,descr,sequence)=FT.getNext(fasta)

        # parse result
        (done,skip,errmsg,stype)=FT.parseResult(stype,eof,label,descr,sequence,None,labels,verbose)
        if errmsg is not None:
            OT.print_error(errmsg)
            nerror=nerror+1
            continue
        if skip:
            nskip=nskip+1
            continue

        slen = len(sequence)
        mlen = slen
        if masked is None:
            masked=FT.getMasked(stype)
        if stype is not None:
            mlen = FT.statCharSequence(sequence,masked)

        # print full sequence length per label
        if perlbl:
            print(('%s %d')%(label,slen))
        sumlen = sumlen + float(slen)
        summlen = summlen + float(mlen)

        if maxlen is None or slen>maxlen:
            maxlen=slen
            maxid=label

        if minlen is None or slen<minlen:
            minlen=slen
            minid=label

    nok = ntot-nskip-nerror
    if nok>0:
        avelen  = sumlen/float(nok)
        avemlen = summlen/float(nok)
    else:
        avelen  = 0.0
        avemlen = 0.0

    if nok==0:
        minlen=0
        minid=''
        maxlen=0
        maxid=''
    if dosh:
        print(('NTOT=%d')%(ntot))
        print(('NERROR=%d')%(nerror))
        print(('NSKIP=%d')%(nskip))
        print(('MINLEN=%d')%(minlen))
        print(('MINLENID=\"%s\"')%(minid))
        print(('MAXLEN=%d')%(maxlen))
        print(('MAXLENID=\"%s\"')%(maxid))
        print(('AVELEN=%f')%(avelen))
        print(('AVEMLEN=%f')%(avemlen))
    else:
        print(('#NTOT:%d')%(ntot))
        print(('#NERROR:%d')%(nerror))
        print(('#NSKIP:%d')%(nskip))
        print(('#MINLEN:%d')%(minlen))
        print(('#MINLENID:%s')%(minid))
        print(('#MAXLEN:%d')%(maxlen))
        print(('#MAXLENID:%s')%(maxid))
        print(('#AVELEN:%f')%(avelen))
        print(('#AVEMLEN:%f')%(avemlen))
        print(('#TOTLEN:%d')%(int(sumlen)))

def cleanupFasta( fsname, trgname, fstype ):
    ec=0
    if trgname is None:
        trgname=fsname
    try:
        # input data
        f = open(fsname)

        # temporary file as target
        trg=Tmp.NamedTemporaryFile()

        # cleanup
        ec = FT.cleanupFasta( f, trg.file, fstype )
        f.close()

        # copy temp file to target
        if ec==0:
            trg.file.flush()
            shutil.copy( trg.name, trgname )
        # close temp file - will also remove the file
        trg.close()

    except IOError as err:
        print("({0}) I/O error: {1}".format(fsname,err))
        ec=255
    #
    if ec>0:
        if ec==1: OT.print_error(('(%s) problem found with fasta file - missing label')%fsname )
        if ec==2: OT.print_error(('(%s) problem found with fasta file - bad fasta type %s')%(fsname,fstype) )
        if ec==3: OT.print_error(('(%s) problem found with fasta file - unable to establish masked character')%fsname )
        if ec>3: OT.print_error(('(%s) problem found with fasta file - return ec = %d')%(fsname,ec) )
    return ec


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Check fasta files')

    parser.add_argument('-f', '--fasta', dest='filename', help='fasta file', metavar='FILE', required=True)
    parser.add_argument('-o', '--fout', dest='fout', help='fasta file output (--clean)', metavar='FILE', default=None)
    parser.add_argument('--expected', dest='expected', help='expected fasta type (with --check)', metavar='string', default=None)
    parser.add_argument('--type',  dest='typefs', help='get data type', action='store_true', default=False )

    parser.add_argument('--label',  dest='label', help='get species label', action='store_true', default=False )

    parser.add_argument('--spsep', dest='spsep', help='species separator', metavar='char', default=None)
    parser.add_argument('--check', dest='check',   help='check type and labels (if spsep is set) of fasta file', action='store_true', default=False )
#    parser.add_argument('--chklabel', dest='chklabel',   help='do also check on labels that they are of format <str1>:<str2>', action='store_true', default=False )
    parser.add_argument('--clean', dest='clean',   help='clean protein fasta file from stop codons', action='store_true', default=False )
    parser.add_argument('--stats', dest='stats', help='output statistics', action='store_true', default=False )
    parser.add_argument('--perlbl', dest='perlbl', help='output statistics per label', action='store_true', default=False )

    parser.add_argument('--title',   dest='title',   help='print title', action='store_true', default=False )
    parser.add_argument('--notype',   dest='notype',   help='ignore checking fasta type', action='store_true', default=False )
    parser.add_argument('--verbose', dest='verbose', help='more output', action='store_true', default=False )

    options = parser.parse_args()

    if not options.check and options.clean:
        OT.print_warning('option --clean is ignored since --check is not set')
        options.clean=False

    # open fasta file and determine the type
    # it will ignore stop codons etc
    fasta = open(options.filename)#,errors='replace')
    found=True
    typefs=None
    dostats = options.stats or options.perlbl
    
    if not options.notype:
        found,typefs = FT.getTypeFile(fasta,options.verbose)

    if not found:
        OT.print_error(('(%s) no valid fasta entry')%(options.filename))
        OT.print_error('for more info, re-run using bin/check_fasta.py --check -f %s --verbose'%options.filename)
        sys.exit(250)

    if options.typefs:
        print(('%s')%(typefs))

    elif options.label:
        found,label = FT.getSpeciesFile(fasta,options.spsep)
        if not found:
            OT.print_error(('(%s) no valid fasta entry or no label found with separator \'%s\'')%(options.filename,options.spsep))
            sys.exit(251)
        print(('%s')%(label))

    elif options.check:
        errcode=0
        if options.expected is None:
            expected=typefs
        else:
            expected=options.expected

        # check fasta if an expected type has been found
        errcode,label,typefsFound=FT.checkFastaType( fasta, expected, False, options.spsep, options.verbose )

        if errcode==0 and options.clean:
            # bad characters found, potentially stop codons
            # clean up, if allowed
            ec=cleanupFasta( options.filename, options.fout, typefs )
            if not ec==0:
                errcode=8
                OT.print_error(('(%s) cleanupFasta returned with error code %d')%(options.filename,ec))

        # check error code
        if errcode==1:
            OT.print_error(('(%s) invalid label found (%s)')%(options.filename,label))
        elif errcode==2:
            OT.print_error(('(%s) ') % (options.filename))
        elif errcode==3:
            OT.print_error(('(%s) inconsistent sequence (%s) not having expected type %s')%(options.filename,label,expected))
        elif errcode==4:
            OT.print_error(('(%s) type found (%s) is not as expected (%s)') % (options.filename,typefsFound,expected))
        elif errcode==5:
            OT.print_error(('(%s) no data in file') % (options.filename))
        elif errcode==6:
            OT.print_error(('(%s) unable to determine fasta type)') % (options.filename))
        elif errcode>7:
            OT.print_error(('(%s) unexpected error code = %d => check with developer!')%(options.filename,errcode))

        if errcode>0:
            fasta.close()
            sys.exit(errcode)

    elif dostats:
        fasta = open(options.filename)#,errors='replace')
        getFastaStats(fasta,False,options.perlbl,options.verbose)
        fasta.close()
