import cv2
import ezdxf
import numpy as np


def draw_hatch(img, entity, color, mask):
    for poly_path in entity.paths.paths:
        # print(poly_path.path_type_flags)
        polygon = np.array([vertex[:-1] for vertex in poly_path.vertices]).astype(int)
        if poly_path.path_type_flags & 1 == 1:
            cv2.fillPoly(img, [polygon], color)
            cv2.fillPoly(mask, [polygon], (255, 255, 255))
        else:
            cv2.fillPoly(img, [polygon], (255, 255, 255))
    return color


def draw_line(img, entity, color, mask):
    p1 = entity.dxf.start[:-1]
    p2 = entity.dxf.end[:-1]
    cv2.line(img, (int(p1[0]), int(p1[1])), (int(p2[0]), int(p2[1])), color, 1)
    cv2.line(mask, (int(p1[0]), int(p1[1])), (int(p2[0]), int(p2[1])), (255, 255, 255), 2)
    return color


def draw_lwpolyline(img, entity, color, mask):
    polyline = []
    a = np.array(entity.lwpoints.values).astype(int)
    while len(a) > 0:
        polyline.append((a[0], a[1]))
        a = a[5:]
    cv2.polylines(img, [np.array(polyline)], entity.closed, color, 1)
    cv2.polylines(mask, [np.array(polyline)], entity.closed, (255, 255, 255), 2)
    return color


def draw_arc(img, entity, color, mask):
    s = entity.dxf.start_angle * np.pi / 180
    e = entity.dxf.end_angle * np.pi / 180
    if s > e:
        s -= 2 * np.pi
    d = (e - s) / (int((e - s) * 180 / np.pi) + 1)
    r = entity.dxf.radius
    cx, cy = entity.dxf.center.xyz[:-1]
    angles = np.arange(s, e + d / 2, d)
    x = cx + r * np.cos(angles)
    y = cy + r * np.sin(angles)
    points = np.column_stack((x, y)).astype(int)
    cv2.polylines(img, [points], abs(s - e) < 1e-9, color, 1)
    cv2.polylines(mask, [points], abs(s - e) < 1e-9, (255, 255, 255), 2)
    return color


def draw_circle(img, entity, color, mask):
    r = entity.dxf.radius
    cx, cy = entity.dxf.center.xyz[:-1]
    cv2.circle(img, (int(cx), int(cy)), int(r), color, 1)
    cv2.circle(mask, (int(cx), int(cy)), int(r), (255, 255, 255), -1)
    return color


def draw_ellipse(img, entity, color, mask):
    cx, cy = entity.dxf.center.xyz[:-1]
    ma = entity.dxf.major_axis.magnitude
    angle = entity.dxf.major_axis.angle_deg
    mi = ma * entity.dxf.ratio
    s = entity.dxf.start_param * 180 / np.pi
    e = entity.dxf.end_param * 180 / np.pi
    if entity.dxf.extrusion.z == -1:
        s = 360 - s
        e = 360 - e
    cv2.ellipse(img, (int(cx), int(cy)), (int(ma), int(mi)), angle, s, e, color, 1)
    cv2.ellipse(mask, (int(cx), int(cy)), (int(ma), int(mi)), angle, s, e, (255, 255, 255), 1)
    return color


def draw_point(img, entity, color, mask):
    cx, cy = entity.dxf.location.xyz[:-1]
    cv2.circle(img, (int(cx), int(cy)), 0, color, 1)
    cv2.circle(mask, (int(cx), int(cy)), 0, (255, 255, 255), -1)
    return color


draw_map = {
    'HATCH': draw_hatch,
    'LINE': draw_line,
    'LWPOLYLINE': draw_lwpolyline,
    'ARC': draw_arc,
    'CIRCLE': draw_circle,
    'ELLIPSE': draw_ellipse,
    'POINT': draw_point,
}


def paint(in_path, out_path, config):
    doc = ezdxf.readfile(in_path)
    extmax, extmin = doc.header['$EXTMAX'], doc.header['$EXTMIN']
    xmin, ymin = np.floor(extmin[:-1]).astype(int)
    xmax, ymax = np.ceil(extmax[:-1]).astype(int)
    img = np.ones((ymax + ymin, xmax + xmin, 3), np.uint8) * 255
    mask = np.zeros_like(img)
    msp = doc.modelspace()
    layers = config.get('layers', {})
    colors = config.get('colors', {})
    # print(doc.layers.entries.keys())
    for layer_name, names in layers.items():
        color = tuple(colors.get(layer_name, [0, 0, 0]))
        for name in names:
            if name not in doc.layers:
                continue
            entities = msp.query('*[layer=="%s"]' % name)
            tmp = np.zeros((ymax + ymin, xmax + xmin), np.uint8)
            for entity in entities:
                if entity.DXFTYPE in draw_map:
                    draw_map[entity.DXFTYPE](img, entity, color, tmp)
                else:
                    print("%s: %s" % (name, entity.DXFTYPE))
            contours, hierarchy = cv2.findContours(tmp, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            cv2.drawContours(mask, contours, -1, color, -1)

    res, img_png = cv2.imencode('.png', cv2.flip(img, 0))
    res, mask_png = cv2.imencode('.png', cv2.flip(mask, 0))
    with open(out_path, 'wb') as f:
        f.write(img_png.tobytes())
    with open(out_path[:-4] + "_mask.png", 'wb') as f:
        f.write(mask_png.tobytes())

[((45, 13, 45, 39), 'numpy.arange', 'np.arange', ({(45, 23, 45, 24): 's', (45, 26, 45, 35): 'e + d / 2', (45, 37, 45, 38): 'd'}, {}), '(s, e + d / 2, d)', True, 'import numpy as np\n'), ((96, 10, 96, 33), 'ezdxf.readfile', 'ezdxf.readfile', ({(96, 25, 96, 32): 'in_path'}, {}), '(in_path)', False, 'import ezdxf\n'), ((101, 11, 101, 29), 'numpy.zeros_like', 'np.zeros_like', ({(101, 25, 101, 28): 'img'}, {}), '(img)', True, 'import numpy as np\n'), ((100, 10, 100, 58), 'numpy.ones', 'np.ones', ({(100, 18, 100, 47): '(ymax + ymin, xmax + xmin, 3)', (100, 49, 100, 57): 'np.uint8'}, {}), '((ymax + ymin, xmax + xmin, 3), np.uint8)', True, 'import numpy as np\n'), ((121, 40, 121, 56), 'cv2.flip', 'cv2.flip', ({(121, 49, 121, 52): 'img', (121, 54, 121, 55): '0'}, {}), '(img, 0)', False, 'import cv2\n'), ((122, 41, 122, 58), 'cv2.flip', 'cv2.flip', ({(122, 50, 122, 54): 'mask', (122, 56, 122, 57): '0'}, {}), '(mask, 0)', False, 'import cv2\n'), ((11, 12, 11, 47), 'cv2.fillPoly', 'cv2.fillPoly', ({(11, 25, 11, 28): 'img', (11, 30, 11, 39): '[polygon]', (11, 41, 11, 46): 'color'}, {}), '(img, [polygon], color)', False, 'import cv2\n'), ((12, 12, 12, 58), 'cv2.fillPoly', 'cv2.fillPoly', ({(12, 25, 12, 29): 'mask', (12, 31, 12, 40): '[polygon]', (12, 42, 12, 57): '(255, 255, 255)'}, {}), '(mask, [polygon], (255, 255, 255))', False, 'import cv2\n'), ((14, 12, 14, 57), 'cv2.fillPoly', 'cv2.fillPoly', ({(14, 25, 14, 28): 'img', (14, 30, 14, 39): '[polygon]', (14, 41, 14, 56): '(255, 255, 255)'}, {}), '(img, [polygon], (255, 255, 255))', False, 'import cv2\n'), ((28, 8, 28, 40), 'numpy.array', 'np.array', ({(28, 17, 28, 39): 'entity.lwpoints.values'}, {}), '(entity.lwpoints.values)', True, 'import numpy as np\n'), ((32, 24, 32, 42), 'numpy.array', 'np.array', ({(32, 33, 32, 41): 'polyline'}, {}), '(polyline)', True, 'import numpy as np\n'), ((33, 25, 33, 43), 'numpy.array', 'np.array', ({(33, 34, 33, 42): 'polyline'}, {}), '(polyline)', True, 'import numpy as np\n'), ((46, 17, 46, 31), 'numpy.cos', 'np.cos', ({(46, 24, 46, 30): 'angles'}, {}), '(angles)', True, 'import numpy as np\n'), ((47, 17, 47, 31), 'numpy.sin', 'np.sin', ({(47, 24, 47, 30): 'angles'}, {}), '(angles)', True, 'import numpy as np\n'), ((48, 13, 48, 36), 'numpy.column_stack', 'np.column_stack', ({(48, 29, 48, 35): '(x, y)'}, {}), '((x, y))', True, 'import numpy as np\n'), ((98, 17, 98, 38), 'numpy.floor', 'np.floor', ({(98, 26, 98, 37): 'extmin[:-1]'}, {}), '(extmin[:-1])', True, 'import numpy as np\n'), ((99, 17, 99, 37), 'numpy.ceil', 'np.ceil', ({(99, 25, 99, 36): 'extmax[:-1]'}, {}), '(extmax[:-1])', True, 'import numpy as np\n'), ((112, 18, 112, 64), 'numpy.zeros', 'np.zeros', ({(112, 27, 112, 53): '(ymax + ymin, xmax + xmin)', (112, 55, 112, 63): 'np.uint8'}, {}), '((ymax + ymin, xmax + xmin), np.uint8)', True, 'import numpy as np\n'), ((118, 34, 118, 99), 'cv2.findContours', 'cv2.findContours', ({(118, 51, 118, 54): 'tmp', (118, 56, 118, 73): 'cv2.RETR_EXTERNAL', (118, 75, 118, 98): 'cv2.CHAIN_APPROX_SIMPLE'}, {}), '(tmp, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)', False, 'import cv2\n'), ((119, 12, 119, 59), 'cv2.drawContours', 'cv2.drawContours', ({(119, 29, 119, 33): 'mask', (119, 35, 119, 43): 'contours', (119, 45, 119, 47): '(-1)', (119, 49, 119, 54): 'color', (119, 56, 119, 58): '(-1)'}, {}), '(mask, contours, -1, color, -1)', False, 'import cv2\n'), ((9, 18, 9, 74), 'numpy.array', 'np.array', ({(9, 27, 9, 73): '[vertex[:-1] for vertex in poly_path.vertices]'}, {}), '([vertex[:-1] for vertex in poly_path.vertices])', True, 'import numpy as np\n')]