"""
Measure the precision and coverage of a fragment library stored in Rosetta
NNmake format.
"""

import os
import multiprocessing
import matplotlib

import csb.apps

import csb.bio.io.wwpdb  as wwpdb
import csb.bio.structure as structure
import csb.bio.fragments
import csb.bio.fragments.rosetta as rosetta
import csb.core
import csb.io.plots


class ExitCodes(csb.apps.ExitCodes):
    
    IO_ERROR = 2
    INVALID_DATA = 3

class AppRunner(csb.apps.AppRunner):
    
    @property
    def target(self):
        return PrecisionApp
     
    def command_line(self):
        
        cpu = multiprocessing.cpu_count()
        cmd = csb.apps.ArgHandler(self.program, __doc__)

        cmd.add_scalar_option('pdb', 'p', str, 'the PDB database (a directory containing all PDB files)', required=True)
        cmd.add_scalar_option('native', 'n', str, 'native structure of the target (PDB file)', required=True)
        cmd.add_scalar_option('chain', 'c', str, 'chain identifier (if not specified, the first chain)', default=None)

        cmd.add_scalar_option('top', 't', int, 'read top N fragments per position', default=25)
        cmd.add_scalar_option('cpu', 'C', int, 'maximum degree of parallelism', default=cpu)
        cmd.add_scalar_option('rmsd', 'r', float, 'RMSD cutoff for precision and coverage', default=1.5)         
        cmd.add_scalar_option('output', 'o', str, 'output directory', default='.')
        
        cmd.add_boolean_option('save-structures', 's', 'create a PDB file for each fragment, superimposed over the native', default=False)
        
        cmd.add_positional_argument('library', str, 'Fragment library file in Rosetta NNmake format')
                        
        return cmd

class PrecisionApp(csb.apps.Application):
    
    def main(self):
        
        if not os.path.isdir(self.args.output):
            PrecisionApp.exit('Output directory does not exist', code=ExitCodes.INVALID_DATA, usage=True)
        
        for file in [self.args.native, self.args.library]:
            if not os.path.isfile(file):
                PrecisionApp.exit('File not found: {0}'.format(file), code=ExitCodes.INVALID_DATA, usage=True)            
        
        self.log('\nLoading {0.library} (top {0.top} per position)... '.format(self.args), ending='')
        
        try:
            library = rosetta.RosettaFragmentMap.read(self.args.library, top=self.args.top)
            self.log('{0} fragments'.format(len(library)))
            
            try:
                native = wwpdb.RegularStructureParser(self.args.native).parse_structure()
                native.accession = native.accession[:4]
                if self.args.chain:
                    native = native.chains[self.args.chain]
                else:
                    native = native.first_chain
            except (structure.Broken3DStructureError, wwpdb.PDBParseError) as se:
                raise ArgumentError("Can't parse native structure: " + str(se))
            
            except (csb.core.ItemNotFoundError) as ce:
                raise ArgumentError("Chain {0!s} not found in native structure".format(ce))            

            self.log('Superimposing fragments...\n')
            si = LibrarySuperimposer(native, library, self.args.pdb, self.args.output,
                                      save=self.args.save_structures, cutoff=self.args.rmsd)
            matches = si.run(self.args.cpu)
            
            info = si.plot(matches)
            self.log(' RMSD Cutoff    {0:>6.2f} A'.format(self.args.rmsd))            
            self.log(' AVG Precision  {0.precision:>6.2f} %'.format(info))
            self.log(' Coverage       {0.coverage:>6.2f} %'.format(info))
            self.log(' RW  Precision  {0.figure}'.format(si))


            
        except ArgumentIOError as ioe:
            PrecisionApp.exit(str(ioe), code=ExitCodes.INVALID_DATA)
        
        except ArgumentError as ae:
            PrecisionApp.exit(str(ae), code=ExitCodes.INVALID_DATA)

        self.log('\nDone.')      

class ArgumentError(ValueError):
    pass

class ArgumentIOError(ArgumentError):
    pass         

class GlobalInfo(object):
    
    def __init__(self, precision, coverage):
        self.precision = precision
        self.coverage = coverage
        
    def __str__(self):
        return '{0.precision:6.2f} {0.coverage:6.2f}'.format(self)
    
    def __sub__(self, rhs):
        return GlobalInfo(self.precision - rhs.precision, self.coverage - rhs.coverage)
        
class LibrarySuperimposer(object):
    
    def __init__(self, native, library, pdb, output, save=False, cutoff=1.5):
    
        if not isinstance(native, structure.Chain):
            raise TypeError(native)
        elif native.length < 1:
            raise ArgumentError('The native chain has no residues')
        
        if not isinstance(library, rosetta.RosettaFragmentMap):
            raise TypeError(library)

        for i in [pdb, output]:
            if not os.path.isdir(i):
                raise ArgumentIOError("{0} is not a valid directory".format(i))

        self._pdb = pdb
        self._native = native
        self._library = library
        self._output = os.path.abspath(output)
        self._tab = os.path.join(self._output, native.entry_id + '.fragments.tab')
        self._figure = os.path.join(self._output, native.entry_id + '.precision.pdf')
        self._out = open(self._tab, 'w')
        self._save = bool(save)
        self._cutoff = float(cutoff)

    def __del__(self):
        try:
            if self._out and not self._out.closed:
                self._out.close()
        except:
            pass
        
    @property
    def table(self):
        return self._tab
    
    @property
    def figure(self):
        return self._figure
            
    def run(self, cpu=1):
        
        tasks = []
        matches = []
        save = None
        if self._save:
            save = self._output
            
        pool = multiprocessing.Pool(cpu)
        
        for source in self._library.sources:
            fragments = self._library.fromsource(source)
            task = pool.apply_async(rmsd, [self._native, source, fragments, self._pdb, save])
            tasks.append(task)
        
        for task in tasks:
            for match in task.get():
                if isinstance(match, csb.core.string):       # error
                    self._out.write(' [E] ')                    
                    self._out.write(match.rstrip())
                    self._out.write('\n')
                else:
                    line = '{0.id}\t{0.qstart}\t{0.qend}\t{0.length}\t{0.rmsd}\n'.format(match)
                    self._out.write(line)
                    matches.append(match)
        
        pool.terminate()
        return matches
                    
    def plot(self, matches):
        
        residues = range(1, self._native.length + 1)
        precision = []
        precision2 = []
        background = []
        covered = set()
        
        all = {}
        positive = {}
        
        for rank in residues:
            all[rank] = 0
            positive[rank] = 0
        
        for match in matches:
            for rank in range(match.qstart, match.qend + 1):                
                all[rank] += 1
                
                if match.rmsd <= self._cutoff:    
                    positive[rank] += 1
                    covered.add(rank)
        
        assert sorted(all.keys()) == sorted(positive.keys()) == residues
        
        for rank in residues:
            if all[rank] == 0:
                precision2.append(0)
                background.append(0)
            else:
                p = positive[rank] * 100.0 / all[rank]
                precision.append(p)
                precision2.append(p)
                background.append(100)
                
        coverage = len(covered) * 100.0 / len(residues)
        if len(precision) > 0:
            avg_precision = sum(precision) / len(precision)
        else:
            avg_precision = 0                 
        
        with csb.io.plots.Chart() as chart:
                        
            chart.plot.bar(residues, background, color='#f5f5f5', linewidth=None, edgecolor='#f5f5f5')
            chart.plot.bar(residues, precision2, color='#5ba9da', linewidth=None, edgecolor='#5ba9da')            

            chart.plot.set_title(self._native.entry_id)
            chart.plot.set_xlabel('Residue')
            chart.plot.set_xlim(1, len(residues))
            chart.plot.set_ylabel('Precision, %')
            chart.plot.set_ylim(0, 100)

            xaxis = chart.plot.axes.xaxis
            yaxis = chart.plot.axes.yaxis            
            #xaxis.set_minor_locator(matplotlib.ticker.IndexLocator(1, 0))
            xaxis.set_major_locator(matplotlib.ticker.IndexLocator(10, 0))   
            xaxis.tick_bottom()
            yaxis.tick_left()
            for t in xaxis.get_major_ticks(): 
                t.tick1On = False 
                t.tick2On = False
            for t in xaxis.get_ticklabels():
                t.set_fontsize(16)
            for t in yaxis.get_ticklabels():
                t.set_fontsize(16)
                                    
            chart.plot.spines["right"].set_visible(False)
            chart.plot.spines["top"].set_visible(False)                             
                        
            try:
                chart.width = 15
                chart.height = 5.5
                
                chart.save(self._figure, chart.formats.PDF)
                
            except IOError as io:
                raise ArgumentIOError("Can't save figure: " + str(io))
        
        return GlobalInfo(avg_precision, coverage)
            
def rmsd(target, source, fragments, pdb, save=None):

    matches = []
    
    try:
        src_file = wwpdb.find(source, [pdb])
        if src_file is None:
            raise IOError("Can't find structure {0} in {1}".format(source, pdb))

        src_structure = wwpdb.RegularStructureParser(src_file).parse_structure()

    except (IOError, structure.Broken3DStructureError, wwpdb.PDBParseError) as ex:
        matches.append('Error parsing {0:5}: {1!s}'.format(source, ex))
        return matches
            
    for fn, fragment in enumerate(fragments):            
            try:                
                if fragment.chain not in ('_', '', None):
                    src_chain = src_structure.chains[fragment.chain]
                else:
                    src_chain = src_structure.first_chain

                try: 
                    query = target.subregion(fragment.qstart, fragment.qend)
                    subject = src_chain.subregion(fragment.start, fragment.end, clone=True)
                    
                except IndexError:
                    matches.append('Fragment {1.source_id:>5} {0.start:>4}-{0.end:>4} is out of range'.format(fragment))                 
                    continue                    
            
                si = query.align(subject)
                match = csb.bio.fragments.FragmentMatch(fragment.id, qstart=fragment.qstart, qend=fragment.qend,
                                                        probability=None, rmsd=si.rmsd, tm_score=None, qlength=target.length)              
                matches.append(match)
                
                if save:
                    dummy = structure.Structure(subject.entry_id)
                    dummy.chains.append(subject)
                    filename = '{0.qstart:}-{0.qend}-{0.source_id}.{1.entry_id}{2}.frag'.format(fragment, query, fn or '')
                    dummy.to_pdb(os.path.join(save, filename))
                    
            except (structure.Broken3DStructureError, IOError) as ex:
                matches.append("Can't superimpose fragment {0}: {1!s}".format(fragment.id, ex))          
                continue
                
    return matches         


def main():
    AppRunner().run()
    
    
if __name__ == '__main__':
    main()