#!/usr/bin/env python 

'''
License
    This file is part of foam-extend.

    foam-extend 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.

    foam-extend 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 foam-extend.  If not, see <http://www.gnu.org/licenses/>.

Author
    Matthias Rauter matthias.rauter@uibk.ac.at
'''


# these 3 lines make python2 compatible with python3
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import argparse
import time
import sys
import shapefile as shp
import re
import os

parser = argparse.ArgumentParser(description='')

parser.add_argument('-case', type=str, help='OpenFOAM case folder', default='.')
parser.add_argument('-fields', type=str, help='input all fieldname (separated by space)')
parser.add_argument('-asPoints', help='export points in area centers instead of cells', action="store_true")
parser.add_argument('-o', type=str, help='output destination')
parser.add_argument('-offsetx', type=float, help='offset mesh by x m', default=0)
parser.add_argument('-offsety', type=float, help='offset mesh by y m', default=0)

args = parser.parse_args()


# reading FA Mesh

print('reading FV mesh from "{}"'.format(args.case+'/constant/polyMesh'))

dat = open(args.case+'/constant/polyMesh/points', 'r')

points = []
while True:
    line=dat.readline()
    if not line or line.startswith('('):
        break
while True:
    line=dat.readline()
    if not line or line.startswith(')'):
        break
    else:
        line = re.search(r'\((.*)\)', line).group(0)[1:-1]
        coords = [float(i) for i in line.split(' ')]
        points.append(coords)

dat = open(args.case+'/constant/faMesh/faceLabels', 'r')

print('reading FA mesh from "{}"'.format(args.case+'/constant/faMesh'))

faceLabels = []
while True:
    line=dat.readline()
    if not line or line.startswith('('):
        break
while True:
    line=dat.readline()
    if not line or line.startswith(')'):
        break
    else:
        faceLabels.append(int(line))

dat = open(args.case+'/constant/polyMesh/faces', 'r')

faces = []
while True:
    line=dat.readline()
    if not line or line.startswith('('):
        break
while True:
    line=dat.readline()
    if not line or line.startswith(')'):
        break
    else:
        lineComplete = False
        while not lineComplete:
            try:
                line = re.search(r'\((.*)\)', line).group(0)[1:-1]
                lineComplete = True
            except:
                line = line.rstrip() + ' ' + dat.readline()
        nodes = [int(i) for i in line.split(' ') if i != '']
        faces.append(nodes)


print('{} total faces'.format(len(faces)))
faces = [faces[i] for i in faceLabels]
print('{} boundary faces'.format(len(faces)))

print('{} total points'.format(len(points)))

vtoa = [-1 for i in points]
atov = []

i = 0
for f in faces:
    for n in f:
        if vtoa[n] == -1:
            vtoa[n] = len(atov)
            atov.append(n)

points = [points[i] for i in atov]

for i in range(len(faces)):
    for j in range(len(faces[i])):
        faces[i][j] = vtoa[faces[i][j]]


print('{} boundary points'.format(len(points)))

for i in range(len(points)):
    points[i][0] -= args.offsetx
    points[i][1] -= args.offsety
rpath = args.case

fields = []
fieldNames = []
fieldNamesAll = args.fields.split(' ')
for i in fieldNamesAll:
    print('reading {}'.format(rpath+'/'+i))
    try:
        dat = open(rpath+'/'+i, 'r')
    except:
        print('unable to open "{}". Skipping...'.format(i))
        continue
    uniform = False
    while True:
        line=dat.readline()
        if not line:
            break
        if 'areaVectorField' in line:
            vector = True
            fields.append([])
            fields.append([])
            fields.append([])
            fields.append([])
            fieldNames.append(i+'_x')
            fieldNames.append(i+'_y')
            fieldNames.append(i+'_z')
            fieldNames.append(i+'_mean')
        if 'areaScalarField' in line:
            vector = False
            fields.append([])
            fieldNames.append(i)
        if line.startswith('internalField'):
            if not 'nonuniform' in line:
                uniform = True
                if vector:
                    uniformValue = re.search(r'\((.*)\)', line).group(0)[1:-1]
                    uniformValue = [float(i) for i in uniformValue.split(' ')]
                    for i in faces:
                        fields[-4].append(uniformValue[0])
                        fields[-3].append(uniformValue[1])
                        fields[-2].append(uniformValue[2])  
                        fields[-1].append((uniformValue[0]**2+uniformValue[1]**2+uniformValue[2]**2**0.5))     
                else:
                    uniformValue = float(re.search(r'[-+]?\d*\.\d+|\d+', line).group(0))
                    for i in faces:
                        fields[-1].append(uniformValue)        
            else:
                line=dat.readline(); line=dat.readline();
            break
    if not uniform:
        while True:
            line = dat.readline()
            if not line or line.startswith(')'):
                break
            if vector:
                v = re.search(r'\((.*)\)', line).group(0)[1:-1]
                v = [float(i) for i in v.split(' ')]
                fields[-4].append(v[0])
                fields[-3].append(v[1])
                fields[-2].append(v[2])
                fields[-1].append((v[0]**2+v[1]**2+v[2]**2)**0.5)
            else:
                fields[-1].append(float(line))

if args.asPoints:
    w = shp.Writer(shp.POINT)
else:
    w = shp.Writer(shp.POLYGON)

if args.asPoints:
    faceCenters_x = []
    faceCenters_y = []
    for f in faces:
        x = 0
        y = 0
        for n in f:
            x += points[n][0]
            y += points[n][1]
        faceCenters_x.append(x/len(f))
        faceCenters_y.append(y/len(f))
        w.point(faceCenters_x[-1], faceCenters_y[-1])

else:
    for face in faces:
    	lines = []
        for nodes in face:
           lines.append([points[nodes][0], points[nodes][1]])
        lines.append([points[face[0]][0], points[face[0]][1]])
    	w.poly(parts=[lines])

    dat.close()

w.field('id','C')

for i in range(len(fieldNames)):
	w.field(fieldNames[i], 'F')

for i in range(len(faces)):
    w.record(i, *tuple([fields[ii][i] for ii in range(len(fieldNames))]))

print('Writing to file "{}"'.format(args.o))
w.save(args.o)

