#!/usr/bin/env python

VERSION = "$Id: doubleGrin.py,v 1.10 2008/08/20 22:17:45 pletzer Exp $"

"""
Compute the response matrices lambda = A * d lambda/dn for 2 concentric
contours

"""

from math import pi, cos, sin, sqrt, floor
from Triangulate import Triangulate
from ellipt2d import ellipt2d
import vector
import dsupralu
import Tkinter
import tkplot
import re
import sys
import copy
import numpy


class doubleGrin:

    def __init__(self, filename, args):
        """
        Constructor

        @param nTor toroidal mode number (>=1)
        """
        self.contours = {}
        
        self.f = file(filename,'w')
        print >>self.f, "# generated by '%s'" % args
        print >>self.f, "# version %s" % VERSION
        import time
        print >>self.f, '# on %s' % time.asctime()

    def readGrinData(self, filename='grin.dat'):
        """
        Read data file containing contour coordinates

        @param filename
        """
        patHeader = re.compile(r'^\s*(\w\d)\s+(\d+)')
        parData   = re.compile(r'^\s+\-?\d+\.\d+[Ee][\+\-]\d+')
        coordName = ''
        coordSize = 0
        for line in file(filename, 'r').readlines():
           
           m = re.match(patHeader, line)
           if m:
               coordName = m.group(1)
               coordSize = m.group(2)
           elif re.match(parData,line) and coordName and (coordSize > 0):
               self.contours[coordName] = self.contours.get(coordName, []) + \
                                          eval('[' + re.sub(r'(\d)\s+([\d\-])','\\1,\\2', line) + ']')

        # checks
        for k in self.contours:
            if k[0] == 'x':
                xName = k
                nx = len(self.contours[xName])
                yName = re.sub(r'x','y', xName)
                ny = len(self.contours[yName])
                if nx != ny:
                    print 'readGrinData ERROR: lengths of %s (%d) and %s (%d) do not match!' % \
                          (xName, nx, yName, ny)
                    print xName, '=', self.contours[xName]
                    print yName, '=', self.contours[yName]
                    sys.exit(1)
               
    def generateGrids(self, boxRatio=0.5, maxArea=None, minAngle=27.0):
        """
        Triangulate internal and external domains

        @param maxArea maximum triangle area
        @param minAngle minimum angle of a triangle
        """

        n1 = len(self.contours['x1'])
        self.pts1 = [ (self.contours['x1'][i], self.contours['y1'][i]) for i in range(n1) ]
        self.segs1 = [(i, i+1) for i in range(n1-1)] + [(n1-1, 0)]
        
        n2 = len(self.contours['x2'])
        self.pts2  = [ (self.contours['x2'][i], self.contours['y2'][i]) for i in range(n2) ]
        self.segs2 = [(n1+i, n1+i+1) for i in range(n2-1)] + [(n1+n2-1, n1)]

        if maxArea==None:
            maxSegLength = max( [
                sqrt( (self.pts2[i+1][0]-self.pts2[i][0])**2 + \
                      (self.pts2[i+1][1]-self.pts2[i][1])**2 ) for \
                    i in range(len(self.pts2)-1) ])
            maxArea = 0.5 * maxSegLength**2
            print 'setting maxArea to ', maxArea
        self.maxArea = maxArea

        #
        # internal domain
        #
        
        # average for contour 1 is hole position
        x0 = reduce(lambda x,y:x+y, self.contours['x1']) / float(n1)
        y0 = reduce(lambda x,y:x+y, self.contours['y1']) / float(n1)
        self.holes = [ (x0, y0) ]
        
        triInternal = Triangulate()
        triInternal.set_points(self.pts1 + self.pts2, [i for i in range(n1+n2)])
        # attributes used to locate segments on contours 1 and 2.
        #
        attrs = [(float(i),) for i in range(n1)] + \
                [(float(n1+i),) for i in range(n2)]
        triInternal.set_attributes(attrs)
        triInternal.set_segments(self.segs1 + self.segs2)
        triInternal.set_holes(self.holes)

        # need to check that there are no Steiner points
        self.minAngle = minAngle
        triInternal.triangulate(area=maxArea, mode='pzq%fe' % minAngle)
        self.gridInternal = triInternal.get_nodes()
        #self.gridInternal.plot()
        self.edgesInternal = triInternal.get_edges()
        self.attrsInternal = triInternal.get_attributes()
        self.subSegsInternal = {}

        # build the sub-segment data structure, which associates a macro-segment
        # to a list of subsegments

        # first set keys
        for seg in self.segs1:
            self.subSegsInternal[seg] = []
        for seg in self.segs2:
            self.subSegsInternal[seg] = []
        
        for e in self.edgesInternal:
            
            if e[1] == 0:
                # internal point
                continue

            i, j = e[0]
            # assume segments are indexed counterclockwise
            ai, aj = self.attrsInternal[i][0], self.attrsInternal[j][0]
            if ai > aj:
                ai, aj = self.attrsInternal[j][0], self.attrsInternal[i][0]
            iSegNode = int(floor(ai))
            jSegNode = iSegNode + 1
            if abs(aj - ai) > 1.0:
                # closing segment
                if aj < n1:
                    iSegNode = n1-1
                    jSegNode = 0
                else:
                    iSegNode = n1+n2-1
                    jSegNode = n1
            ijSeg = (iSegNode, jSegNode)

            self.subSegsInternal[ijSeg] += [ ijSeg ]
        #
        # External domain
        #

        minY = min(self.contours['y2'])
        maxY = max(self.contours['y2'])
        Ly = maxY - minY
        minY -= 0.5 * boxRatio * Ly
        maxY += 0.5 * boxRatio * Ly
        minX = 1.e-6
        maxX = max(self.contours['x2'])
        Lx = maxX - minX
        maxX += boxRatio * Lx

        # enforce same number of points as on contour 1
        # so can re-use the list of segments on contour 2
        n3 = n1
        assert(n3 % 4 == 0) # 4 sides
        
        dy = (maxY - minY) / float(n3/4)
        dx = (maxX - minX) / float(n3/4)
        # box boundary points. Take same number as internal
        # contour, which allows us to share contour (2)
        # south
        self.pts3 = [ (minX+i*dx, minY) for i in range(n3/4) ]
        # east
        self.pts3 += [ (maxX, minY+i*dy) for i in range(n3/4) ]
        # north
        self.pts3 += [ (maxX-i*dx, maxY) for i in range(n3/4) ]
        # west
        self.pts3 += [ (minX, maxY-i*dy) for i in range(n3/4) ]
        self.segs3 = [(i, i+1) for i in range(n3-1) ] \
                     + [(n3-1, 0)]

        triExternal = Triangulate()
        triExternal.set_points(self.pts3 + self.pts2)
        # attributes used to locate segments on contours 2.
        #
        attrs = [(-1.0,) for i in range(n3)] + \
                [(float(n3+i),) for i in range(n2)]
        triExternal.set_attributes(attrs)
        triExternal.set_segments(self.segs3  + self.segs2)
        triExternal.set_holes(self.holes)

        triExternal.triangulate(area=maxArea, mode='pzq%fe' % minAngle)
        self.gridExternal = triExternal.get_nodes()
        #self.gridExternal.plot()
        self.edgesExternal = triExternal.get_edges()
        self.attrsExternal = triExternal.get_attributes()
        self.subSegsExternal = {}

        # build the sub-segment data structure, which associates a macro-segment
        # to a list of subsegments

        # first set keys
        for seg in self.segs2:
            self.subSegsExternal[seg] = []
        
        for e in self.edgesExternal:
            
            if e[1] == 0:
                # internal point
                continue

            i, j = e[0]
            # assume segments are indexed counterclockwise
            ai, aj = self.attrsExternal[i][0], self.attrsExternal[j][0]
            if ai < 0.0 or aj < 0.0:
                # this segment is on contour 3
                continue
            if ai > aj:
                ai, aj = self.attrsExternal[j][0], self.attrsExternal[i][0]
            iSegNode = int(floor(ai))
            jSegNode = iSegNode + 1
            if abs(aj - ai) > 1.0:
                # closing segment
                if aj < n1:
                    iSegNode = n1-1
                    jSegNode = 0
                else:
                    iSegNode = n1+n2-1
                    jSegNode = n1
            ijSeg = (iSegNode, jSegNode)

            self.subSegsExternal[ijSeg] += [ ijSeg ]

        self.nt1 = n1
        self.nt2 = n2
        self.nt3 = n3
        
        self.A11 = numpy.mat( numpy.zeros( (self.nt1, self.nt1), numpy.float64 ) )
        self.A12 = numpy.mat( numpy.zeros( (self.nt1, self.nt2), numpy.float64 ) )
        self.A21 = numpy.mat( numpy.zeros( (self.nt2, self.nt1), numpy.float64 ) )
        self.A22 = numpy.mat( numpy.zeros( (self.nt2, self.nt2), numpy.float64 ) )
        self.A33 = numpy.mat( numpy.zeros( (self.nt2, self.nt2), numpy.float64 ) )

    def assembleInternal(self, nTor):

        f = 'x'
        g = '%d/x' % nTor**2
        s = None
        self.equInternal = ellipt2d(self.gridInternal, f, g, s)
        AInternal, self.rhs0Internal = self.equInternal.stiffnessMat()

        # do the lu decomposition once and for all
        m, n = AInternal.size()
        self.handlesInternal = dsupralu.new(AInternal, n)
        colperm = 2
        dsupralu.colperm(self.handlesInternal[0], colperm)
        dsupralu.lu(self.handlesInternal[0])

    def assembleExternal(self, nTor):

        f = 'x'
        g = '%d/x' % nTor**2
        s = None
        self.equExternal = ellipt2d(self.gridExternal, f, g, s)
        AExternal, self.rhs0External = self.equExternal.stiffnessMat()

        # do the lu decomposition once and for all
        m, n = AExternal.size()
        self.handlesExternal = dsupralu.new(AExternal, n)
        colperm = 2
        dsupralu.colperm(self.handlesExternal[0], colperm)
        dsupralu.lu(self.handlesExternal[0])

    def solveInternal(self, ij):
        
        """
        Solve in internal domain for unit d lambda/dn on boundary

        @param ij segment spanning nodes (i,j) (assume i->j is counterclockwise)
        """

        i, j = ij
        value = 1.0
        rhsInternal = copy.copy(self.rhs0Internal)
        # iterate over Steiner points
        for s in self.subSegsInternal[ij]:
            k, l = s
            xk, yk = self.gridInternal.x(k), self.gridInternal.y(k)
            xl, yl = self.gridInternal.x(l), self.gridInternal.y(l)
            # mid point
            x     = (xk + xl)/2.0
            y     = (yk + yl)/2.0
            dx    = xl - xk
            dy    = yl - yk
            ds    = sqrt( dx**2 + dy**2 )
            # split contribution
            contrb = x * value * ds /2.
            rhsInternal[k] += contrb
            rhsInternal[l] += contrb
            
        # solve
        self.chiInternal = vector.vector(
            dsupralu.solve(self.handlesInternal[0], rhsInternal))

    def solveExternal(self, ij):
        
        """
        Solve in external domain for unit d lambda/dn on boundary

        @param ij segment spanning nodes (i,j) (assume i->j is counterclockwise)
        """

        i, j = ij
        value = 1.0 # normal derivative
        rhsExternal = copy.copy(self.rhs0External)
        # iterate over Steiner points
        for s in self.subSegsExternal[ij]:
            k, l = s
            xk, yk = self.gridExternal.x(k), self.gridExternal.y(k)
            xl, yl = self.gridExternal.x(l), self.gridExternal.y(l)
            # mid point
            x     = (xk + xl)/2.0
            y     = (yk + yl)/2.0
            dx    = xl - xk
            dy    = yl - yk
            ds    = sqrt( dx**2 + dy**2 )
            # split contribution
            contrb = x * value * ds /2.
            rhsExternal[k] += contrb
            rhsExternal[l] += contrb
            
        # solve
        self.chiExternal = vector.vector(
            dsupralu.solve(self.handlesExternal[0], rhsExternal))

    def computeInternalMatrices(self):
        """
        """
        
        # sources on contour 1
        for ij in self.segs1:
            self.solveInternal(ij)
            self.A11[:, ij[0]] = self.getChi1()
            self.A21[:, ij[0]] = self.getInternalChi2()

        # sources on contour 2
        for ij in self.segs2:
            self.solveInternal(ij)
            self.A12[:, ij[0] - self.nt1] = self.getChi1()
            self.A22[:, ij[0] - self.nt1] = self.getInternalChi2()

    def computeExternalMatrices(self):
        """
        """
        
        # sources on contour 2
        for ij in self.segs2:
            self.solveExternal(ij)
            self.A33[:, ij[0] - self.nt3] = self.getExternalChi2()

##         if plot:
##             WIDTH = 600
##             HEIGHT = 600
##             root = Tkinter.Tk()
##             frame = Tkinter.Frame(root)
##             frame.pack()
##             canvas = Tkinter.Canvas(bg="White",width=WIDTH,height=HEIGHT)
##             canvas.pack()
##             tkplot.tkplot(canvas,self.gridInternal, self.chiInternal, 0,0,1,
##                           title="nTor=%d"%self.nTor, 
##                           WIDTH=WIDTH,HEIGHT=HEIGHT)
##             root.mainloop()
        
    def getChi1(self):
        """
        Get solution on contour 1

        @return solution vector
        """

        return numpy.matrix([self.chiInternal[i] for i in range(self.nt1)]).T

    def getInternalChi2(self):
        """
        Get internal solution on contour 2

        @return solution vector
        """

        return numpy.matrix( \
            [self.chiInternal[i] for i in range(self.nt1, self.nt1+self.nt2)] \
            ).T

    def getExternalChi2(self):
        """
        Get external solution on contour 2

        @return solution vector
        """

        return numpy.matrix( \
            [self.chiExternal[i] for i in range(self.nt3, self.nt3+self.nt2)] \
            ).T

    def writeData(self, nTor):
        """
        Save matrices in file
        """
        nCols = 8
        for kind in '11', '12', '21', '22', '33':
            data = eval('self.A' + kind)
            print >>self.f,'name=%s nTor=%d sizes=[%d,%d] ' % \
              ('A'+kind, nTor, data.shape[0], data.shape[1]),
            descr = 'response matrix lambda(nod%s)/d lambda(seg%s)/dn' \
                    % (kind[0], kind[1])
            print >>self.f,'descr="%s"' % descr
            nElems = len(data.flat)
            nRows = nElems // nCols
            for ir in range(nRows):
                print >>self.f, ('%12e ')*nCols % \
                  tuple(data.flat[ir*nCols:(ir+1)*nCols].flat)
            nRemd = nElems - nRows*nCols
            # last line
            print >>self.f, ('%12e ')*nRemd  % \
                  tuple(data.flat[nRows*nCols:].flat)
            

    def testResponse(self, nTor, gradLambd1, delLambd32):
        """
        Test response by solving complete problem in one sweep
        """
        # using response matrices, determine lambda(1) given
        # d lambda/d n and the jump delLambda (=lambdaExt - lambdaInt)

        from numpy.linalg import inv
        
        bMat = inv(self.A33 - self.A22)
        lambd1_a = (self.A11 + self.A12 * bMat * self.A21) * gradLambd1 + \
                 self.A12 * bMat * delLambd32

        from matplotlib import pylab
        pylab.plot(lambd1_a.flat)

        # now do the same solving the whole problem
        tr = Triangulate()
        tr.set_points(self.pts1 + self.pts2 + self.pts3)
        atts = [ (1.0,) for i in range(self.nt1) ] + \
               [ (2.0,) for i in range(self.nt2) ] + \
               [ (3.0,) for i in range(self.nt3) ]
        tr.set_attributes(atts)
        segs = self.segs1 + \
               [ (i, i+1) for i in range(self.nt1+self.nt2, self.nt1+self.nt2+self.nt3-1) ] + \
               [ (self.nt1+self.nt2+self.nt3-1, self.nt1+self.nt2) ] 
        tr.set_segments(segs)
        tr.set_holes( self.holes )
        tr.triangulate(area=self.maxArea, mode='pzq%fe' % self.minAngle)
        grd = tr.get_nodes()
        grd.plot()
        edgs = tr.get_edges()
        atts = tr.get_attributes()
        f = 'x'
        g = '%d/x' % nTor**2
        s = None
        equ = ellipt2d(grd, f, g, s)
        A, rhs0 = equ.stiffnessMat()
        rhs = copy.copy(rhs0)
        # do the lu decomposition once and for all
        m, n = A.size()
        handles = dsupralu.new(A, n)
        colperm = 2
        dsupralu.colperm(handles[0], colperm)
        dsupralu.lu(handles[0])
        # apply Neumann BCs
        n1 = len(self.pts1)
        for e in edgs:
            onBoundary = e[1]
            if onBoundary == 1.0:
                i, j = e[0]
                print '***** i,j = ', i,j, atts[i], atts[j]

                if atts[i][0]==1.0 and atts[j][0]==1.0:
                    print '++++++++++++++++++++++++++++++++++++++++++++++++++'
                    # on contour 1
                    value = 1.0 # for now !!!!!!!
                    xi, yi = grd.x(i), grd.y(i)
                    xj, yj = grd.x(j), grd.y(j)
                    x = 0.5*(xi + xj)
                    y = 0.5*(yi + yj)
                    dx, dy = xj - xi, yj - yi
                    ds = sqrt( dx**2 + dy**2 )
                    contrib = x * value * ds /2.
                    rhs[i] += contrib
                    rhs[j] += contrib
        print rhs
        sol = dsupralu.solve(handles[0], rhs)
        lambd1_b = [sol[i] for i in range(self.nt1)]

        pylab.hold(True)
        print lambd1_b
        pylab.plot(lambd1_b, 'r-')
        pylab.show()
        
    
##############################################################################
def main():
    from optparse import OptionParser
    import sys

    parser = OptionParser(version=VERSION)
    
    parser.add_option('-n', '--nTor', action='store', type="int",
                      dest="nTor", help="Toroidal mode number (!=0)",
                      default=1)
    parser.add_option('-f', '--file', action='store', type="string",
                      dest="filename", default="grin.dat",
                      help="Grin filename describing contours",
                      )
    parser.add_option('-o', '--output', action='store', type="string",
                      dest="outFile", default="grinOut.dat",
                      help="Output file containing response matrices",
                      )
    parser.add_option('-p', '--plot', action='store_true',
                      dest="plot", help="To plot solution",
                      default=False)

    options, args = parser.parse_args(sys.argv)
    
    dl = doubleGrin(filename=options.outFile, args=sys.argv)
    dl.readGrinData(filename=options.filename)
    dl.generateGrids(minAngle=27.0)
    dl.assembleInternal(nTor=options.nTor)
    dl.computeInternalMatrices()
    dl.assembleExternal(nTor=options.nTor)
    dl.computeExternalMatrices()
    dl.writeData(nTor=options.nTor)

    gradLambd1 = numpy.matrix( numpy.ones( ( len(dl.contours['x1']), ), numpy.float64 ) ).T
    delLambd32 = numpy.matrix( numpy.zeros( ( len(dl.contours['x2']), ), numpy.float64 ) ).T
    dl.testResponse(options.nTor, gradLambd1, delLambd32)

    print dl.A11
    

if __name__ == '__main__': main()
