#!/usr/bin/env python3
  
  '''
  Convert GEDI-type bounding boxes to CSV format
  '''
  
  import logging
  import os
  import sys
  import time
  import glob
  import csv
  import imghdr
  from PIL import Image
  import argparse
  import pdb
  import cv2
  import numpy as np
  import xml.etree.ElementTree as ET
  
  sin = np.sin
  cos = np.cos
  pi = np.pi
  
  def Rotate2D(pts, cnt, ang=90):
      M = np.array([[cos(ang),-sin(ang)],[sin(ang),cos(ang)]])
      res = np.dot(pts-cnt,M)+cnt
      return M, res
  
  def npbox2string(npar):
      if np.shape(npar)[0] != 1:
          print('Error during CSV conversion
  ')
      c1,r1 = npar[0][0],npar[0][1]
      c2,r2 = npar[0][2],npar[0][3]
      c3,r3 = npar[0][4],npar[0][5]
      c4,r4 = npar[0][6],npar[0][7]
  
      return c1,r1,c2,r2,c3,r3,c4,r4
  
  # cv2.minAreaRect() returns a Box2D structure which contains following detals - ( center (x,y), (width, height), angle of rotation )
  # Get 4 corners of the rectangle using cv2.boxPoints()
  class GEDI2CSV(object):
      ''' Initialize the extractor'''
      def __init__(self, logger, args):
          self._logger = logger
          self._args = args
  
      '''
      Segment image with GEDI bounding box information
      '''
      def csvfile(self, coords, polys, baseName, pgrot):
  
          ''' for writing the files '''
          writePath = self._args.outputDir
          writePath = os.path.join(writePath,'')
          if os.path.isdir(writePath) != True:
              os.makedirs(writePath)
  
          rotlist = []
  
          header=['ID','name','col1','row1','col2','row2','col3','row3','col4','row4','confidence','truth','pgrot','bbrot','qual','script','text_type']
          conf=100
          write_ctr = 0
          if len(coords) == 0 and len(polys) == 0:
              self._logger.info('Found %s with no text content',(baseName))
              print('...Found %s with no text content' % (baseName))
              return
              
          strPos = writePath + baseName
  
          ''' for each group of coordinates '''
          for i in coords:
  
              [id,x,y,w,h,degrees,text,qual,script,text_type] = i
              contour = np.array([(x,y),(x+w,y),(x+w,y+h),(x,y+h)])
              """First rotate around upper left corner based on orientationD keyword"""
              M, rot = Rotate2D(contour, np.array([x,y]), degrees*pi/180)
              rot = np.int0(rot)
  
              # rot is the 8 points rotated by degrees
              # pgrot is the rotation after extraction, so save
              # save rotated points to list or array
              rot = np.reshape(rot,(-1,1)).T
              c1,r1,c2,r2,c3,r3,c4,r4 = npbox2string(rot)
              
              bbrot = degrees
              rotlist.append([id,baseName + '_' + id + '.png',c1,r1,c2,r2,c3,r3,c4,r4,conf,text,pgrot,bbrot,qual,script,text_type])
  
          # if there are polygons, first save the text
          for j in polys:
              arr = []
              [id,poly_val,text,qual,script,text_type] = j
              for i in poly_val:
                  arr.append(eval(i))
  
              contour = np.asarray(arr)
              convex = cv2.convexHull(contour)
              rect = cv2.minAreaRect(convex)
              box = cv2.boxPoints(rect)
              box = np.int0(box)
              box = np.reshape(box,(-1,1)).T
              c1,r1,c2,r2,c3,r3,c4,r4 = npbox2string(box)
              
              bbrot = 0.0
              rotlist.append([id,baseName + '_' + id + '.png',c1,r1,c2,r2,c3,r3,c4,r4,conf,text,pgrot,bbrot,qual,script,text_type])
          # then write out all of list to file
          with open(strPos + ".csv", "w", encoding="utf-8") as f:
              writer = csv.writer(f)
              writer.writerow(header)
              for row in rotlist:
                  writer.writerow(row)
                  write_ctr += 1
          return write_ctr
      
  
  def main(args):
      startTime = time.clock()
      writePath = args.outputDir
      if os.path.isdir(writePath) != True:
          os.makedirs(writePath)
      ''' Setup logging '''
      logger = logging.getLogger(__name__)
      logger.setLevel(logging.INFO)
      if args.log:
          handler = logging.FileHandler(args.log)
          handler.setLevel(logging.INFO)
          formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
          handler.setFormatter(formatter)
          logger.addHandler(handler)
      gtconverter = GEDI2CSV(logger, args)
      namespaces = {"gedi" : "http://lamp.cfar.umd.edu/media/projects/GEDI/"}
      keyCnt=0
      fileCnt = 0
      line_write_ctr = 0
      line_error_ctr = 0
      
      '''
      Get all XML files in the directory and sub folders
      '''
      for root, dirnames, filenames in os.walk(args.inputDir, followlinks=True):
          for file in filenames:
              if file.lower().endswith('.xml'):
                  fullName = os.path.join(root,file)
                  baseName = os.path.splitext(fullName)
                  fileCnt += 1
                  ''' read the XML file '''
                  tree = ET.parse(fullName)
                  gedi_root = tree.getroot()
                  child = gedi_root.findall('gedi:DL_DOCUMENT',namespaces)[0]
                  totalpages = int(child.attrib['NrOfPages'])
                  coordinates=[]
                  polygons = []
                  if args.ftype == 'boxed':
                      fileTypeStr = 'col'
                  elif args.ftype == 'transcribed':
                      fileTypeStr = 'Text_Content'
                  else:
                      print('Filetype must be either boxed or transcribed!')
                      logger.info('Filetype must be either boxed or transcribed!')
                      sys.exit(-1)
                  
                  if args.quality == 'both':
                      qualset = {'Regular','Low-Quality'}
                  elif args.quality == 'low':
                      qualset = {'Low-Quality'}
                  elif args.quality == 'regular':
                      qualset = {'Regular'}
                  else:
                      print('Quality must be both, low or regular!')
                      logger.info('Quality must be both, low or regular!')
                      sys.exit(-1)
                      
                      
  
                  ''' and for each page '''
                  for i, pgs in enumerate(child.iterfind('gedi:DL_PAGE',namespaces)):
                          
                      if 'GEDI_orientation' not in pgs.attrib:
                          pageRot=0
                      else:
                          pageRot = int(pgs.attrib['GEDI_orientation'])
                          logger.info(' PAGE ROTATION %s, %s' % (fullName, str(pageRot)))
  
                      ''' find children for each page '''
                      for zone in pgs.findall('gedi:DL_ZONE',namespaces):
  
                          if zone.attrib['gedi_type']=='Text' and zone.attrib['Type'] in \
                              ('Machine_Print','Confusable_Allograph','Handwriting') and zone.attrib['Quality'] in qualset:
                              if zone.get('polygon'):
                                  keyCnt+=1
                                  polygons.append([zone.attrib['id'],zone.get('polygon').split(';'),
                                                   zone.get('Text_Content'),zone.get('Quality'),zone.get('Script'),zone.get('Type')])
                              elif zone.get(fileTypeStr) != None:
                                  keyCnt+=1
                                  coord = [zone.attrib['id'],int(zone.attrib['col']),int(zone.attrib['row']),
                                                      int(zone.attrib['width']), int(zone.attrib['height']),
                                                      float(zone.get('orientationD',0.0)),
                                                      zone.get('Text_Content'),zone.get('Quality'),zone.get('Script'),zone.get('Type')]
                                  coordinates.append(coord)
  
                  if len(coordinates) > 0 or len(polygons) > 0:
                      line_write_ctr += gtconverter.csvfile(coordinates, polygons, os.path.splitext(file)[0], pageRot)
                  else:
                      print('...%s has no applicable content' % (baseName[0]))
  
      print('complete...total files %d, lines written %d' % (fileCnt, line_write_ctr))
  
  
  ''' Args and defaults '''
  def parse_arguments(argv):
      parser = argparse.ArgumentParser()
  
      parser.add_argument('--inputDir', type=str, help='Input directory', required=True)
      parser.add_argument('--outputDir', type=str, help='Output directory', required=True)
      parser.add_argument('--ftype', type=str, help='GEDI file type (either "boxed" or "transcribed")', default='transcribed')
      parser.add_argument('--quality', type=str, help='GEDI file quality (either "both" or "low" or "regular")', default='regular')
      parser.add_argument('--log', type=str, help='Log directory', default='./GEDI2CSV_enriched.log')
  
      return parser.parse_args(argv)
  
  ''' Run '''
  if __name__ == '__main__':
      main(parse_arguments(sys.argv[1:]))