# Copyright (c) 2020-2022, Manfred Moitzi
# License: MIT License
from typing import cast
import pathlib
import math
import random
import numpy as np

import ezdxf
from ezdxf import zoom
from ezdxf.enums import TextEntityAlignment
from ezdxf.math import Vec3, Matrix44, Z_AXIS, Y_AXIS, X_AXIS
from ezdxf.entities import Circle, Arc, Ellipse, Insert, Text, MText, Hatch

CWD = pathlib.Path("~/Desktop/Outbox").expanduser()
if not CWD.exists():
    CWD = pathlib.Path(".")

UNIFORM_SCALING = [
    (-1, 1, 1),
    (1, -1, 1),
    (1, 1, -1),
    (-2, -2, 2),
    (2, -2, -2),
    (-2, 2, -2),
    (-3, -3, -3),
]
NON_UNIFORM_SCALING = [
    (-1, 2, 3.1),
    (1, -2, 3.2),
    (1, 2, -3.3),
    (-3.4, -2, 1),
    (3.5, -2, -1),
    (-3.6, 2, -1),
    (-3.7, -2, -1),
]

SCALING_WITHOUT_REFLEXION = [(2, 2, 2), (1, 2, 3)]


def setup_csys_blk(name: str):
    blk = doc.blocks.new(name)
    blk.add_line((0, 0, 0), X_AXIS, dxfattribs={"color": 1})
    blk.add_line((0, 0, 0), Y_AXIS, dxfattribs={"color": 3})
    blk.add_line((0, 0, 0), Z_AXIS, dxfattribs={"color": 5})


def random_angle():
    return random.uniform(0, math.tau)


def synced_scaling(
    entity, chk, axis_vertices=None, sx: float = 1, sy: float = 1, sz: float = 1
):
    entity = entity.copy()
    entity.scale(sx, sy, sz)
    m = Matrix44.scale(sx, sy, sz)
    chk = list(m.transform_vertices(chk))
    if axis_vertices:
        axis_vertices = list(m.transform_vertices(axis_vertices))
        return entity, chk, axis_vertices
    return entity, chk


def synced_rotation(
    entity, chk, axis_vertices=None, axis=Z_AXIS, angle: float = 0
):
    entity = entity.copy()
    entity.rotate_axis(axis, angle)
    m = Matrix44.axis_rotate(axis, angle)
    chk = list(m.transform_vertices(chk))
    if axis_vertices:
        axis_vertices = list(m.transform_vertices(axis_vertices))
        return entity, chk, axis_vertices
    return entity, chk


def synced_translation(
    entity, chk, axis_vertices=None, dx: float = 0, dy: float = 0, dz: float = 0
):
    entity = entity.copy()
    entity.translate(dx, dy, dz)
    m = Matrix44.translate(dx, dy, dz)
    chk = list(m.transform_vertices(chk))
    if axis_vertices:
        axis_vertices = list(m.transform_vertices(axis_vertices))
        return entity, chk, axis_vertices
    return entity, chk


def synced_transformation(entity, chk, m: Matrix44):
    entity = entity.copy()
    entity.transform(m)
    chk = list(m.transform_vertices(chk))
    return entity, chk


def add(msp, entity, vertices, layer="0"):
    entity.dxf.layer = layer
    entity.dxf.color = 2
    msp.entitydb.add(entity)
    msp.add_entity(entity)
    msp.add_polyline3d(vertices, dxfattribs={"layer": "vertices", "color": 6})


def circle(radius=1, count=16):
    circle_ = Circle.new(
        dxfattribs={"center": (0, 0, 0), "radius": radius}, doc=doc
    )
    control_vertices = list(circle_.vertices(np.linspace(0, 360, count)))
    return circle_, control_vertices


def arc(radius=1, start=30, end=150, count=8):
    arc_ = Arc.new(
        dxfattribs={
            "center": (0, 0, 0),
            "radius": radius,
            "start_angle": start,
            "end_angle": end,
        },
        doc=doc,
    )
    control_vertices = list(arc_.vertices(arc_.angles(count)))
    return arc_, control_vertices


def ellipse(
    major_axis=(1, 0),
    ratio: float = 0.5,
    start: float = 0,
    end: float = math.tau,
    count: int = 8,
):
    major_axis = Vec3(major_axis).replace(z=0)
    ellipse_ = Ellipse.new(
        dxfattribs={
            "center": (0, 0, 0),
            "major_axis": major_axis,
            "ratio": min(max(ratio, 1e-6), 1),
            "start_param": start,
            "end_param": end,
        },
        doc=doc,
    )
    control_vertices = list(ellipse_.vertices(ellipse_.params(count)))
    axis_vertices = list(
        ellipse_.vertices([0, math.pi / 2, math.pi, math.pi * 1.5])
    )
    return ellipse_, control_vertices, axis_vertices


def insert():
    return (
        Insert.new(
            dxfattribs={
                "name": "UCS",
                "insert": (0, 0, 0),
                "xscale": 1,
                "yscale": 1,
                "zscale": 1,
                "rotation": 0,
                "layer": "insert",
            },
            doc=doc,
        ),
        [(0, 0, 0), X_AXIS, Y_AXIS, Z_AXIS],
    )


def main_ellipse(layout):
    entity, vertices, axis_vertices = ellipse(
        start=math.pi / 2, end=-math.pi / 2
    )
    axis = Vec3.random()
    angle = random_angle()
    entity, vertices, axis_vertices = synced_rotation(
        entity, vertices, axis_vertices, axis=axis, angle=angle
    )
    entity, vertices, axis_vertices = synced_translation(
        entity,
        vertices,
        axis_vertices,
        dx=random.uniform(-2, 2),
        dy=random.uniform(-2, 2),
        dz=random.uniform(-2, 2),
    )

    for sx, sy, sz in UNIFORM_SCALING + NON_UNIFORM_SCALING:
        entity0, vertices0, axis0 = synced_scaling(
            entity, vertices, axis_vertices, sx, sy, sz
        )
        add(layout, entity0, vertices0, layer=f"new ellipse")
        layout.add_line(
            axis0[0],
            axis0[2],
            dxfattribs={"color": 6, "linetype": "DASHED", "layer": "old axis"},
        )
        layout.add_line(
            axis0[1],
            axis0[3],
            dxfattribs={"color": 6, "linetype": "DASHED", "layer": "old axis"},
        )
        p = list(entity0.vertices([0, math.pi / 2, math.pi, math.pi * 1.5]))
        layout.add_line(
            p[0], p[2], dxfattribs={"color": 1, "layer": "new axis"}
        )
        layout.add_line(
            p[1], p[3], dxfattribs={"color": 3, "layer": "new axis"}
        )


def main_multi_ellipse(layout):
    m = Matrix44.chain(
        Matrix44.scale(1.1, 1.3, 1),
        Matrix44.z_rotate(math.radians(10)),
        Matrix44.translate(1, 1, 0),
    )
    entity, vertices, axis_vertices = ellipse(
        start=math.pi / 2, end=-math.pi / 2
    )

    for index in range(5):
        entity, vertices = synced_transformation(entity, vertices, m)
        add(layout, entity, vertices)


def main_insert(layout):
    entity, vertices = insert()
    entity, vertices = synced_translation(entity, vertices, dx=1, dy=0, dz=0)
    axis = Vec3.random()
    angle = random_angle()

    for sx, sy, sz in NON_UNIFORM_SCALING:
        # 1. scale
        entity0, vertices0 = synced_scaling(
            entity, vertices, sx=sx, sy=sy, sz=sz
        )
        # 2. rotate
        entity0, vertices0 = synced_rotation(
            entity0, vertices0, axis=axis, angle=angle
        )
        # 3. translate
        entity0, vertices0 = synced_translation(
            entity0,
            vertices0,
            dx=random.uniform(-2, 2),
            dy=random.uniform(-2, 2),
            dz=random.uniform(-2, 2),
        )
        layout.entitydb.add(entity0)
        layout.add_entity(entity0)
        origin, x, y, z = list(vertices0)
        layout.add_line(origin, x, dxfattribs={"color": 2, "layer": "new axis"})
        layout.add_line(origin, y, dxfattribs={"color": 4, "layer": "new axis"})
        layout.add_line(origin, z, dxfattribs={"color": 6, "layer": "new axis"})

        for line in entity0.virtual_entities():
            line.dxf.layer = "exploded axis"
            line.dxf.color = 7
            layout.entitydb.add(line)
            layout.add_entity(line)


def main_insert2(layout):
    entity, vertices = insert()
    m = Matrix44.chain(
        Matrix44.scale(-1.1, 1.1, 1),
        Matrix44.z_rotate(math.radians(10)),
        Matrix44.translate(1, 1, 1),
    )
    doc.layers.new("exploded axis", dxfattribs={"color": -7})

    for i in range(5):
        entity, vertices = synced_transformation(entity, vertices, m)
        layout.entitydb.add(entity)
        layout.add_entity(entity)

        origin, x, y, z = list(vertices)
        layout.add_line(origin, x, dxfattribs={"color": 2, "layer": "new axis"})
        layout.add_line(origin, y, dxfattribs={"color": 4, "layer": "new axis"})
        layout.add_line(origin, z, dxfattribs={"color": 6, "layer": "new axis"})

        for line in entity.virtual_entities():
            line.dxf.layer = "exploded axis"
            line.dxf.color = 7
            layout.entitydb.add(line)
            layout.add_entity(line)


def main_text(layout):
    content = "{}RSKNZQ"

    def text(num):
        height = 1.0
        width = 1.0
        p1 = Vec3(0, 0, 0)

        t = Text.new(
            dxfattribs={
                "text": content.format(
                    num
                ),  # should easily show reflexion errors
                "height": height,
                "width": width,
                "rotation": 0,
                "layer": "text",
            },
            doc=doc,
        )
        t.set_placement(p1, align=TextEntityAlignment.LEFT)
        tlen = height * len(t.dxf.text) * width
        p2 = p1.replace(x=tlen)
        p3 = p2.replace(y=height)
        p4 = p1.replace(y=height)
        v = [p1, p2, p3, p4, p3.lerp(p4), p2.lerp(p3)]
        return t, v

    def add_box(vertices):
        p1, p2, p3, p4, center_top, center_right = vertices
        layout.add_line(p1, p2, dxfattribs={"color": 1, "layer": "rect"})
        layout.add_line(p2, p3, dxfattribs={"color": 3, "layer": "rect"})
        layout.add_line(p3, p4, dxfattribs={"color": 1, "layer": "rect"})
        layout.add_line(p4, p1, dxfattribs={"color": 3, "layer": "rect"})
        layout.add_line(
            center_right, p1, dxfattribs={"color": 2, "layer": "rect"}
        )
        layout.add_line(
            center_right, p4, dxfattribs={"color": 2, "layer": "rect"}
        )
        layout.add_line(
            center_top, p1, dxfattribs={"color": 4, "layer": "rect"}
        )
        layout.add_line(
            center_top, p2, dxfattribs={"color": 4, "layer": "rect"}
        )

    entity0, vertices0 = text(1)
    entity0, vertices0 = synced_rotation(
        entity0, vertices0, axis=Z_AXIS, angle=math.radians(30)
    )
    entity0, vertices0 = synced_translation(entity0, vertices0, dx=3, dy=3)

    for i, reflexion in enumerate([(1, 2), (-1, 2), (-1, -2), (1, -2)]):
        rx, ry = reflexion
        m = Matrix44.chain(
            Matrix44.scale(rx, ry, 1),
        )
        entity, vertices = synced_transformation(entity0, vertices0, m)
        entity.dxf.text = content.format(i + 1)

        layout.add_entity(entity)
        add_box(vertices)


def main_mtext(layout):
    content = "{}RSKNZQ"

    def mtext(num):
        height = 1.0
        width = 1.0
        p1 = Vec3(0, 0, 0)

        t = MText.new(
            dxfattribs={
                "char_height": height,
                "width": width,
                "text_direction": (1, 0, 0),
                "attachment_point": 7,
                "layer": "text",
            },
            doc=doc,
        )
        t.text = content.format(num)
        tlen = height * len(t.text) * width
        p2 = p1.replace(x=tlen)
        p3 = p2.replace(y=height)
        p4 = p1.replace(y=height)
        v = [p1, p2, p3, p4, p3.lerp(p4), p2.lerp(p3)]
        return t, v

    def add_box(vertices):
        p1, p2, p3, p4, center_top, center_right = vertices
        layout.add_line(p1, p2, dxfattribs={"color": 1, "layer": "rect"})
        layout.add_line(p2, p3, dxfattribs={"color": 3, "layer": "rect"})
        layout.add_line(p3, p4, dxfattribs={"color": 1, "layer": "rect"})
        layout.add_line(p4, p1, dxfattribs={"color": 3, "layer": "rect"})
        layout.add_line(
            center_right, p1, dxfattribs={"color": 2, "layer": "rect"}
        )
        layout.add_line(
            center_right, p4, dxfattribs={"color": 2, "layer": "rect"}
        )
        layout.add_line(
            center_top, p1, dxfattribs={"color": 4, "layer": "rect"}
        )
        layout.add_line(
            center_top, p2, dxfattribs={"color": 4, "layer": "rect"}
        )

    entity0, vertices0 = mtext(1)
    entity0, vertices0 = synced_rotation(
        entity0, vertices0, axis=Z_AXIS, angle=math.radians(30)
    )
    entity0, vertices0 = synced_translation(entity0, vertices0, dx=3, dy=3)

    for i, reflexion in enumerate([(1, 2), (-1, 2), (-1, -2), (1, -2)]):
        rx, ry = reflexion
        m = Matrix44.chain(
            Matrix44.scale(rx, ry, 1),
        )
        entity, vertices = synced_transformation(entity0, vertices0, m)
        entity.text = content.format(i + 1)

        layout.add_entity(entity)
        add_box(vertices)


def hatch_polyline(msp, edge_path=True):
    vertices = [(0, 0, 1), (10, 0), (10, 10, -0.5), (0, 10)]
    hatch = msp.add_hatch(color=1)
    hatch.paths.add_polyline_path(vertices, is_closed=1)
    if edge_path:
        hatch.paths.arc_edges_to_ellipse_edges()
    lwpoly = msp.add_lwpolyline(
        vertices, format="xyb", close=True, dxfattribs={"color": 1}
    )
    return hatch, lwpoly


def main_uniform_hatch_polyline(layout):
    entitydb = layout.doc.entitydb
    hatch, lwpolyline = hatch_polyline(layout, edge_path=False)
    m = Matrix44.chain(
        Matrix44.scale(-1.1, 1.1, 1),
        Matrix44.z_rotate(math.radians(10)),
        Matrix44.translate(1, 1, 1),
    )
    for index in range(4):
        color = 2 + index

        hatch = hatch.copy()
        entitydb.add(hatch)
        hatch.dxf.color = color
        hatch.transform(m)

        lwpolyline = lwpolyline.copy()
        entitydb.add(lwpolyline)
        lwpolyline.dxf.color = color
        lwpolyline.transform(m)

        layout.add_entity(lwpolyline)
        layout.add_entity(hatch)


def main_non_uniform_hatch_polyline(layout, spline=False):
    entitydb = layout.doc.entitydb
    hatch, lwpolyline = hatch_polyline(layout)
    if spline:
        hatch.paths.all_to_spline_edges()

    m = Matrix44.chain(
        Matrix44.scale(-1.1, 1.1, 1),
        Matrix44.z_rotate(math.radians(10)),
        Matrix44.translate(1, 1, 1),
    )
    for index in range(4):
        color = 2 + index
        hatch = hatch.copy()
        entitydb.add(hatch)
        hatch.dxf.color = color
        hatch.transform(m)
        layout.add_entity(hatch)


def main_ellipse_hatch(layout, spline=False):
    def draw_ellipse_axis(ellipse):
        center = ellipse.center
        major_axis = ellipse.major_axis
        msp.add_line(center, center + major_axis)

    entitydb = layout.doc.entitydb
    hatch = cast(Hatch, layout.add_hatch(color=1))
    path = hatch.paths.add_edge_path()
    path.add_line((0, 0), (5, 0))
    path.add_ellipse(
        (2.5, 0), (2.5, 0), ratio=0.5, start_angle=0, end_angle=180, ccw=1
    )
    if spline:
        hatch.paths.all_to_line_edges(spline_factor=4)

    chk_ellipse, chk_vertices, _ = ellipse(
        (2.5, 0), ratio=0.5, start=0, end=math.pi
    )
    chk_ellipse, chk_vertices = synced_translation(
        chk_ellipse, chk_vertices, dx=2.5
    )

    m = Matrix44.chain(
        Matrix44.scale(1.1, 1.3, 1),
        Matrix44.z_rotate(math.radians(15)),
        Matrix44.translate(1, 1, 0),
    )
    for index in range(3):
        color = 2 + index

        hatch = hatch.copy()
        entitydb.add(hatch)
        hatch.dxf.color = color
        hatch.transform(m)
        layout.add_entity(hatch)

        ellipse_edge = hatch.paths[0].edges[1]
        if not spline:
            draw_ellipse_axis(ellipse_edge)

        chk_ellipse, chk_vertices = synced_transformation(
            chk_ellipse, chk_vertices, m
        )
        add(layout, chk_ellipse, chk_vertices)


def add_hatch_for_all_ellipses(layout):
    for ellipse in layout.query("ELLIPSE"):
        hatch = layout.add_hatch(
            color=2,
            dxfattribs={
                "extrusion": ellipse.dxf.extrusion,
                "layer": "HATCH",
            },
        )
        path = hatch.paths.add_edge_path()
        e = ellipse.construction_tool().to_ocs()
        hatch.dxf.elevation = e.center.replace(x=0, y=0)
        edge = path.add_ellipse(
            center=e.center.vec2,
            major_axis=e.major_axis.vec2,
            ratio=e.ratio,
        )
        edge.start_param = e.start_param
        edge.end_param = e.end_param


if __name__ == "__main__":
    doc = ezdxf.new("R2000", setup=True)
    setup_csys_blk("UCS")
    msp = doc.modelspace()
    main_ellipse(msp)
    main_multi_ellipse(msp)
    add_hatch_for_all_ellipses(msp)
    main_text(msp)
    main_mtext(msp)
    main_insert(msp)
    main_insert2(msp)
    main_uniform_hatch_polyline(msp)
    main_ellipse_hatch(msp, spline=True)
    main_non_uniform_hatch_polyline(msp, spline=True)
    zoom.extents(msp)
    doc.saveas(CWD / "transform.dxf")