tikzsimplify

#!/usr/bin/env python3

from __future__ import print_function #Use python3 print statments in python2.7
import re
import argparse
import sys
from math import log10,pow
from operator import lt


tikzpic = re.compile(r'\\begin{tikzpicture}(.*)\end{tikzpicture}',re.MULTILINE|re.DOTALL)
axis = re.compile(r'\\begin{axis}\s*?(\[.*?\])?(.*?)\\end{axis}',re.MULTILINE|re.DOTALL)
plot= re.compile(r'\\addplot\s*?(\[.*?\])?\s*?table\s*?(\[.*?\])?\s*?{(.*?)};',re.MULTILINE|re.DOTALL)
comments = re.compile(r'%.*?$.?',re.MULTILINE|re.DOTALL)

#Match plots of the same style that lead on from each other.
#matlab2tikz sometimes creates these and I don't know why
lead_on_plots = re.compile(r'''( \\addplot \s* (\[.*?\]) \s*
                                   table \s*? (\[.*?\]) \s* { % \s*
                                       (.*?)
                                       ^(.*?)$) \s*};
                              \s*  \\addplot .*? \2 .*? \5
                            ''',re.X|re.MULTILINE|re.DOTALL)


class Heap(object):
    def __init__(self,comp=lt,*args):
        self.array = []
        self.comp = comp

    def push(self,obj):
        i = len(self.array)
        obj.__heappos__ = i
        self.array.append(obj)
        self.up(i)


    def pop(self,rem=None):
        if rem is None:
            rem = self.array[0]

        last = self.array.pop()
        i = rem.__heappos__
        if i != len(self.array):
            #Swap the last element into the removed position then restore heap
            self.array[i] = last
            last.__heappos__ = i
            if self.comp(rem,last):
                self.down(i)
            else:
                self.up(i)
        
        del rem.__heappos__
        return rem

    def update(self,obj):
        # if obj not in self.index:
        #     raise IndexError('object not in heap')
        self.pop(obj)
        self.push(obj)

    def __len__(self):
        return len(self.array)
    def __bool__(self):
        return len(self.array) == 0
    def __getitem__(self,key):
        return self.array[key]

    def up(self,i):
        '''
        Swap the object at index i with its parent until the heap is restored
        '''
        obj = self.array[i]
        while i > 0:
            up = (i-1) >> 1
            parent = self.array[up]
            if self.comp(obj,parent):
                self.array[i] = parent
                parent.__heappos__ = i
                i = up
                continue
            break
        self.array[i] = obj
        obj.__heappos__ = i

    def down(self,i):
        obj = self.array[i]
        n = len(self.array)

        while True:
            right = (i+1)<<1
            left = right-1

            # Find the smallest of current, and two childrend
            down = i
            child = obj
            if left < n and self.comp(self.array[left],child):
                down = left
                child = self.array[left]
            if right < n and self.comp(self.array[right],child):
                down = right
                child = self.array[right]
            #Heap is restored
            if down == i:
                break

            #Swap the moving object and the correct child
            self.array[i] = child
            child.__heappos__ = i

            self.array[down] = obj
            obj.__heappos__ = down
            i = down
            

class Point(object):
    #Doubly linked list of coordinates
    def __init__(self,x,y=None,p=None,n=None):
        
        object.__setattr__(self,'x',x)
        object.__setattr__(self,'y',y)
        object.__setattr__(self,'n',n)
        object.__setattr__(self,'p',p)
        object.__setattr__(self,'area',None)
        if p:
            p.n = self
        if n:
            n.p = self

        self.update_area()

    def __setattr__(self,name,value):
        object.__setattr__(self,name,value)
        if name[0] != '_':
            self.update_area()
        
    def __hash__(self):
        return hash((self.x,self.y))
    def __eq__(self,other):
        # return self.area() == other.area()
        return (self.x == other.x) and (self.y == other.y)
    def __str__(self):
        return '%0.6f,%0.6f' % (self.x,self.y)
    def __repr__(self):
        return str(self.area)

    def update_area(self):
        '''Area this point makes with its two neighbours'''
        if self.n is None or self.p is None:
            return self.area == None
        a = self.p
        b = self
        c = self.n
        ar = abs((b.y - c.y)*a.x + (c.y-a.y)*b.x + (a.y-b.y)*c.x)/2
        self.__dict__['area'] = ar

    def __lt__(a,b):
        return a.area < b.area

    def simplify(self,tol):
        '''
         Perform the Visvalingam--Whyatt algorithm to simplify the path.
         Points are added to a min-heap where the value associated is the 
         area of the triangle made with its two neighbouring points. Points
         are then removed from the list, minimum first, until the smallest
         area is larger than 'tol'.
        
         The first and last points are skipped.
        '''
        h = Heap()
        pt = self.n
        while pt.n:
            h.push(pt)
            pt = pt.n
        i = 0
        while len(h) > 0 and h[0].area < tol:
            i += 1
            pt = h.pop()
            pt.p.n = pt.n
            pt.n.p = pt.p
            #The area of the neighbouring triangles has now changed,
            #so we need to fix the heap
            if pt.p.p is not None:
                h.update(pt.p)

            if pt.n.n is not None:
                h.update(pt.n)

        printv("%d points removed"%i)

def table_to_path(s,logx=False,logy=False):
    rows = s.split('\n')
    head = None
    tail = None

    i = 1
    for row in rows:
        if row:
            x,y = row.split()
            if x == 'nan' or y == 'nan\\\\':
                continue
            x,y = float(x),float(y.strip('\\'))

            if logx:
                x = log10(x)
            if logy:
                y = log10(y)
            tail = Point(x,y,tail)
            if i ==1:
                head = tail
            i += 1
    printv("%d points found"%i)

    return head

def path_to_table(head,logx=False,logy=False):
    '''
    Render points in the path to table format
    '''
    pathlist=[]
    pt = head
    while pt:
        x,y = pt.x,pt.y
        if logx:
            x = pow(10,x)
        if logy:
            y = pow(10,y)
        pathlist.append('\t%0.6g\t%0.6g\\\\' % (x,y))
        pt = pt.n
    return '\n'.join(pathlist)

def simplify_plot(match):
    opts = match.group(1)

    logx= False
    xmode = re.search('xmode=(.*?),',opts)
    if xmode and xmode.group(1) == 'log':
        printv("using logx mode")
        logx=True
    
    logy= False
    ymode = re.search('ymode=(.*?),',opts)
    if ymode and ymode.group(1) == 'log':
        printv("using logy mode")
        logy=True

    xmin = float(re.search('xmin=(.*?),',opts).group(1))
    xmax = float(re.search('xmax=(.*?),',opts).group(1))
    if logx:
        xmin,xmax = log10(xmin),log10(xmax)

    ymin = float(re.search('ymin=(.*?),',opts).group(1))
    ymax = float(re.search('ymax=(.*?),',opts).group(1))
    if logy:
        ymin,ymax = log10(ymin),log10(ymax)

    if tol is None:
        tab_tol = (xmax-xmin)*(ymax-ymin)/(15*9*800)
        printv("selecting tol of: %g"%tab_tol)
    else:
        tab_tol = tol
    
    def simplify_table(match):
        path = table_to_path(comments.sub('',match.group(3)),logx,logy)
        if path is not None:
            path.simplify(tab_tol)
            return '\\addplot %s\n  table%s{%%\n%s\n};' % ( match.group(1) or '', match.group(2) or '', path_to_table(path,logx,logy))


    return  '\\begin{axis}%s\n%s\n\\end{axis}'%( match.group(1) or '', plot.sub(simplify_table,match.group(2)) or '')


def printv(s):
    if verbose:
        print(s,file=sys.stderr)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Simplify matlab2tikz generated plots.')
    parser.add_argument('infile',metavar='infile.tex',nargs=1,help="file to simplify. Use '-' for stdin")
    parser.add_argument('-t',metavar='tol', 
            help='tolerence of simplify algorithm (defaults to best guess)',
            type=float)
    parser.add_argument('-v','--verbose',  action='store_true',help="print info to stderr",default=False)
    parser.add_argument('-i','--in-place',  action='store_true',help="overwrite input file after simplification",default=False)
    parser.add_argument('-o',help='output to outfile.tex rather than stdout',metavar='outfile.tex',default='-')
    args = parser.parse_args()

    verbose = args.verbose
    try:
        tol = args.tol
    except AttributeError:
        tol = None

    if args.infile == '-':
        f = sys.stdin
    else:
        f = open(args.infile,'r')

    s = f.read()
    f.close()
    s = axis.sub(simplify_plot,s)

    # s,n = lead_on_plots.subn(r'\1',s)
    # if n>0:
    #     printv("Lead-on plots detected. ")
    #     printv("I have merged them, but you may want to use  ('maxChunkLength',Inf) in matlab2tikz")
    #     s = axis.sub(simplify_plot,s)

    if args.in_place:
        f = open(args.infile,'w')
    elif 'outfile' not in args or args.outfile == '-':
        f = sys.stdout
    else:
        f = open(args.outfile,'w')

    f.write(s)
    f.close()