#!/usr/bin/env python3
#-----------------------------------------------------------------------------
# Copyright (C) 2017 Correlated Solutions, Inc.
#
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
# 
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
# REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
# AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
# INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
# LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE
# OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
# PERFORMANCE OF THIS SOFTWARE.
#
# This program illustrates how to create project files from a template and
# call Vic-3D in batch mode to analyze them
import argparse
import glob
import os
import subprocess
import sys
import time

import numpy as np

from shutil import which
from VicPy import *

from zipfile import PyZipFile
import lxml.etree as etree

class VicProject:
    root = None
    def init(self):
        self.root = None
        self.images = []
    def load(self, f):
        p = PyZipFile(f, 'r')
        parser = etree.XMLParser(remove_blank_text=True)
        self.root = etree.fromstring(p.read("project.xml"), parser)
        p.close()
    def remove_images(self):
        # Remove images and some project items that require binary data
        rmv = ['//calibration/lcodetarget', '//files/datafile', '//plotdata',
                   '//files/reference', '//files/deformed' ]
        for t in rmv:
            for i in self.root.xpath(t):
                i.getparent().remove(i)
    def save(self, filename, reference_image):
        self.remove_images()
        files = self.root.find('.//files')
        ce = etree.SubElement(files, 'reference')
        ce.text = reference_image
        q = self.root.xpath('//projectaois/aoinode')
        if len(q):
            q[0].attrib['img'] = reference_image
        
        s = etree.tostring(self.root, encoding="utf-8", xml_declaration=True,
                               pretty_print=True, doctype='<!DOCTYPE vpml>').decode('utf-8')
        # Save the xml inside a zip archive
        z = PyZipFile(filename, "w")
        z.writestr("project.xml", s)
        z.close()

def find_vic3d(name):
    if which(name) != None:
        return name
    root = r'C:\Program Files\Correlated Solutions'
    subdirs = [d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))]
    for d in subdirs:
        f = os.path.join(root, d)
        f = os.path.join(f, 'Vic3D.exe')
        if os.path.isfile(f) and os.access(f, os.X_OK):
            return f
    return None

def check_file(f):
    if not os.path.splitext(f)[0].endswith('_0'):
        print('{} is not a camera 0 file, skipping...'.format(f), file=sys.stderr)
        return False
    return (os.path.isfile(f))

__desc__ = '''
This program uses a Vic-3D project file as a template to analyze the shape for a
series of image files. The program calls Vic-3D in batch mode for the shape
analysis and then converts the resulting output file to csv.
'''
    
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__desc__)
    parser.add_argument('template', help='Vic-3D project file used as template.')
    parser.add_argument('image', nargs='+', help='Image file or glob pattern (must be camera 0).')
    parser.add_argument('--exe', '-e', dest='exe', default='Vic3D.exe',
                            help='Vic-3D executable to call')
    args = parser.parse_args()

    exe = find_vic3d(args.exe)
    if exe == None:
        print('Could not find Vic-3D executable.', file=sys.stderr)
        exit(-1)
        
    p = VicProject()
    try:
        p.load(args.template)
    except:
        print('Could not load template from {}'.format(args.template), file=sys.stderr)
        os._exit(-1)
        
    img = []
    for i in args.image:
        if check_file(i):
            img.append(i)
        else:
            for q in glob.glob(i):
                if check_file(q):
                    img.append(q)

    for i in range(len(img)):
        t = time.time()
        # save the project file to a temporary file
        p.save('tmp_project_file.z3d', img[i])
        # call Vic-3D to analyze it
        print('Processing image {}'.format(img[i]))
        subprocess.call([exe, '-L', 'log.txt', '-R', 'tmp_project_file.z3d'])
        # load the resulting output file
        data = VicDataSet()
        base = os.path.splitext(img[i])[0]
        if not data.load(base + '.out'):
            print('Could not load data file {}'.format(df), file=sys.stderr)
            os._exit(-1)
        # extract X, Y, Z coordinates and the sigma value
        val = data.asArray([ 'X', 'Y', 'Z', 'sigma' ])
        # if the sigma value is negative, there is no valid data
        # we use this as a mask to extract valid data only
        mask = val['sigma'] >= 0
        val = np.extract(mask, val)
        P = np.column_stack((val['X'], val['Y'], val['Z']))
        # save a csv file
        np.savetxt(base + '.csv', P, delimiter=',')
        print('Wrote shape data to {}'.format(base + '.csv'))
        try:
            os.remove('tmp_project_file.z3d')
            os.remove(base + '.out')
        except OSError:
            pass
        t = time.time() - t
        print('Points analyzed: {}, time: {:.2f}s'.format(P.shape[0], t))
