"""Extract points from a dxf file, and save in csv format

Example:
    Extract coordinates from all layers
    > python dxf_to_xyz.py file.dxf

    Extract coordinates from one layer
    > python dxf_to_xyz.py file.dxf 'survey'
"""

__author__ = "D. Howe"
__version__ = "0.1.0"
__email__ = "d.howe@wrl.unsw.edu.au"

import os
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import argparse
import dxfgrabber


def point_to_row(point, layer, dxf_type, dxf_id):
    """Create a dict from a point with layer and e properties"""
    row = {key: val for key, val in zip(['x', 'y', 'z'], [*point])}
    row['layer'] = layer
    row['type'] = dxf_type
    row['id'] = dxf_id

    return row


def print_layers(dxf_name):
    dxf = dxfgrabber.readfile(dxf_name)

    # Get data from DTM
    print('Available layers in {}'.format(dxf_name))
    layers = sorted([layer.name for layer in dxf.layers])
    for layer in layers:
        print(layer)


def convert(dxf_name, layer_name):
    # Read dxf file
    dxf = dxfgrabber.readfile(dxf_name)

    if layer_name:
        print("\nExtracting data from layer '{}'".format(layer_name))
        entities = [
            e for e in dxf.entities if e.layer.lower() == layer_name.lower()
        ]
    else:
        # Take all layers if no layer name given
        print('\nExtracing data from all layers')
        entities = [e for e in dxf.entities]

    # Get coordinates of entities
    data = []
    for e in entities:
        if e.dxftype == 'POINT':
            point = e.point
            row = point_to_row(point, e.layer, e.dxftype, e.handle)
            data.append(row)
        elif e.dxftype == 'LINE':
            points = [e.start, e.end]
            for point in points:
                row = point_to_row(point, e.layer, e.dxftype, e.handle)
                data.append(row)
        elif e.dxftype == 'POLYLINE':
            for vertex in e.vertices:
                point = vertex.location
                row = point_to_row(point, e.layer, e.dxftype, e.handle)
                data.append(row)
        elif e.dxftype == 'LWPOLYLINE':
            for point in e.points:
                row = point_to_row(point, e.layer, e.dxftype, e.handle)
                data.append(row)

    # Add to dataframe
    points = pd.DataFrame(data)

    # Check if there are z coordinates
    if 'z' not in points.columns:
        points['z'] = np.nan

    # Reorder columns
    points = points[['layer', 'type', 'id', 'x', 'y', 'z']]

    # Export points
    if layer_name:
        csv_name = dxf_name.replace('.dxf', '-{}.csv'.format(layer_name))
    else:
        csv_name = dxf_name.replace('.dxf', '.csv')
    points.to_csv(csv_name, index=False)


def main():
    # Set up command line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument('dxf_name', help='path to dxf file', default=None)
    parser.add_argument(
        '-n', metavar='LAYER_NAME', help='name of layer', default=None)
    parser.add_argument(
        '-l', action='store_true', help='show list of layers', default=False)
    args = parser.parse_args()

    # Parse arguments
    dxf_name = args.dxf_name
    layer_name = args.n
    show_layers_only = args.l

    if show_layers_only:
        # Print list of layers only
        print_layers(dxf_name)
    else:
        # Convert to csv
        convert(dxf_name, layer_name)


if __name__ == '__main__':
    main()