#!/usr/bin/env python3

import tempfile
import subprocess
import colorsys
import textwrap
import math
from pathlib import Path
import webbrowser
import scipy
import numpy as np
import os

import click


class Tag:
    """ Helper class to ease creation of SVG. All API functions that create SVG allow you to substitute this with your
    own implementation by passing a ``tag`` parameter. """

    def __init__(self, name, children=None, root=False, **attrs):
        if (fill := attrs.get('fill')) and isinstance(fill, tuple):
            attrs['fill'], attrs['fill-opacity'] = fill
        if (stroke := attrs.get('stroke')) and isinstance(stroke, tuple):
            attrs['stroke'], attrs['stroke-opacity'] = stroke
        self.name, self.attrs = name, attrs
        self.children = children or []
        self.root = root

    def __str__(self):
        prefix = '<?xml version="1.0" encoding="utf-8"?>\n' if self.root else ''
        opening = ' '.join([self.name] + [f'{key.replace("__", ":").replace("_", "-")}="{value}"' for key, value in self.attrs.items()])
        if self.children:
            children = '\n'.join(textwrap.indent(str(c), '  ') for c in self.children)
            return f'{prefix}<{opening}>\n{children}\n</{self.name}>'
        else:
            return f'{prefix}<{opening}/>'


    @classmethod
    def setup_svg(kls, tags, bounds, margin=0, unit='mm', pagecolor='white', inkscape=False):
        (min_x, min_y), (max_x, max_y) = bounds

        if margin:
            min_x -= margin
            min_y -= margin
            max_x += margin
            max_y += margin

        w, h = max_x - min_x, max_y - min_y
        w = 1.0 if math.isclose(w, 0.0) else w
        h = 1.0 if math.isclose(h, 0.0) else h

        if inkscape:
            tags.insert(0, kls('sodipodi:namedview', [], id='namedview1', pagecolor=pagecolor,
                    inkscape__document_units=unit))
            namespaces = dict(
                xmlns="http://www.w3.org/2000/svg",
                xmlns__xlink="http://www.w3.org/1999/xlink",
                xmlns__sodipodi='http://sodipodi.sourceforge.net/DTD/sodipodi-0.dtd',
                xmlns__inkscape='http://www.inkscape.org/namespaces/inkscape')

        else:
            namespaces = dict(
                xmlns="http://www.w3.org/2000/svg",
                xmlns__xlink="http://www.w3.org/1999/xlink")

        return kls('svg', tags,
                width=f'{w}{unit}', height=f'{h}{unit}',
                viewBox=f'{min_x} {min_y} {w} {h}',
                style=f'background-color:{pagecolor}',
                **namespaces,
                root=True)



@click.command()
@click.option('--steps', type=int, default=4)
@click.option('--mesh-thickness', type=float, default=1.6)
@click.option('--offset', type=float, default=10)
@click.option('--mesh-space', type=float, default=3)
@click.option('--mesh-width', type=float, default=15)
@click.option('--start', type=click.Choice(['center', 'offset']))
@click.option('--initial-radius', type=float, default=14)
@click.option('--wire-diameter', type=float, default=3, help='Wire diameter for STL export')
@click.argument('out_svg', type=click.Path(dir_okay=False, path_type=Path))
@click.argument('out_stl', type=click.Path(dir_okay=False, path_type=Path), required=False)
def cli(out_svg, out_stl, wire_diameter, mesh_thickness, offset, mesh_space, mesh_width, start, initial_radius, steps):
    tags = []
    dim_tags = []

    hls_svg = lambda h, s, l: '#'+''.join(f'{round(c*255):02x}' for c in colorsys.hls_to_rgb(h/360, l/100, s/100))

    current_color = lambda s, l: hls_svg(0, l, s)
    next_color = lambda s, l: hls_svg(240, l, s)

    # https://stackoverflow.com/questions/13069446/simple-fill-pattern-in-svg-diagonal-hatching
    defs = []
    defs.append(Tag('pattern', [Tag('path', stroke=current_color(40, 70), stroke_width=1,
                        d='M -1,1 l 2,-2 M 0,4 l 4,-4 M 3,5 l2,-2')],
        id='hatch_current', patternUnits='userSpaceOnUse', width='4', height='4', patternTransform='rotate(-45) scale(0.2)'))
    defs.append(Tag('pattern', [Tag('path', stroke=next_color(40, 70), stroke_width=1,
                        d='M -1,1 l 2,-2 M 0,4 l 4,-4 M 3,5 l2,-2')],
        id='hatch_next', patternUnits='userSpaceOnUse', width='4', height='4', patternTransform='scale(0.2)'))
    tags.append(Tag('defs', defs))

    current_fill, next_fill = 'url(#hatch_current)', 'url(#hatch_next)'


    circle = lambda pt, r, c, w=0.2, op=1.0: Tag('circle', stroke=c, fill='none', stroke_width=w, cx=pt[0], cy=pt[1], r=r, stroke_opacity=op)

    current_radius = initial_radius

    def ring(pt, c, fill):
        nonlocal current_radius

        c_line = c(40, 70)
        c_area = fill
        c_bar = c(40, 40)

        r1 = current_radius
        r2 = math.hypot(mesh_width/2, r1 + mesh_thickness)

        tags.append(circle(pt, (r1+r2)/2, c_area, r2-r1, 0.5))
        tags.append(circle(pt, r1, c_line))
        tags.append(circle(pt, r2, c_line))

        px, py = pt
        tags.append(Tag('rect', x=px + r1, y=py-mesh_width/2, width=mesh_thickness, height=mesh_width,
                        fill=c_bar))
        tags.append(Tag('rect', x=px - r1 - mesh_thickness, y=py-mesh_width/2, width=mesh_thickness, height=mesh_width,
                        fill=c_bar))

        current_radius = r2

    def dimension(p1, p2, h=5, line_offset=1, arrow_offset=.5, inner_arrow_min=5, arrow_head_size=1, arrow_length=3,
                  connecting_line_offset=1, color='black', line_width=0.2, text_offset=3, text_offset_max=15,
                  text_color=None, text_format='{:.2f} mm', font_size=2, offset_1=0, offset_2 = 0):
        if text_color is None:
            text_color = color

        x1, y1 = p1
        x2, y2 = p2
        dx, dy = x2-x1, y2-y1
        l = math.hypot(dx, dy)
        nx, ny = -dy/l, dx/l
        ax, ay = dx/l, dy/l

        if l > inner_arrow_min:
            xa1, ya1 = x1 + ax*arrow_offset, y1 + ay*arrow_offset
            xa2, ya2 = x2 - ax*arrow_offset, y2 - ay*arrow_offset
            cl_off = 2*arrow_offset + arrow_length
        else:
            xa2, ya2 = x1 - ax*arrow_offset, y1 - ay*arrow_offset
            xa1, ya1 = x2 + ax*arrow_offset, y2 + ay*arrow_offset
            cl_off = 2*arrow_offset

        angle = math.degrees(math.atan2(dy, dx))

        c_off = h - connecting_line_offset
        cl_off = min(l/2, cl_off)
        yield Tag('path', fill='none', stroke=color, stroke_width=line_width,
                        d=f'''
                        M {x1-nx*line_width/2},{y1-ny*line_width/2} l {nx*line_width},{ny*line_width}
                        M {x1+nx*(line_offset + offset_1)},{y1+ny*(line_offset+offset_1)} l {nx*(h-offset_1)},{ny*(h-offset_1)}
                        M {x2-nx*line_width/2},{y2-ny*line_width/2} l {nx*line_width},{ny*line_width}
                        M {x2+nx*(line_offset+offset_2)},{y2+ny*(line_offset+offset_2)} l {nx*(h-offset_2)},{ny*(h-offset_2)}

                        M {x1+nx*c_off + ax*cl_off},{y1+ny*c_off+ay*cl_off}
                        l {ax*(l-2*cl_off)},{ay*(l-2*cl_off)}
                        ''')

        xa1, ya1 = xa1 + nx*c_off, ya1 + ny*c_off
        xa2, ya2 = xa2 + nx*c_off, ya2 + ny*c_off

        aw, al = arrow_head_size, arrow_head_size*2
        ab = aw*0.2
        for xa, ya, al, arrow_length, abl in [(xa1, ya1, al, arrow_length, ab), (xa2, ya2, -al, -arrow_length, -ab)]:
            yield Tag('path', fill=color, stroke=color, stroke_width=line_width,
                            d=f'''
                            M {xa},{ya}
                            l {nx*aw/2 + ax*al},{ny*aw/2 + ay*al}
                            l {-nx*aw/2 - ax*abl},{-ny*aw/2 - ay*abl}
                            l {ax*(arrow_length-al+abl)},{ay*(arrow_length-al+abl)}
                            l {-ax*(arrow_length-al+abl)},{-ay*(arrow_length-al+abl)}
                            l {-nx*aw/2 + ax*abl},{-ny*aw/2 + ay*abl}
                            Z
                            ''')

        if not -90 <= angle < 90:
            angle = (angle - 180) % 360
            baseline = 'auto'
        else:
            baseline = 'hanging'

        if l > text_offset_max:
            text_offset = 0
            baseline = 'middle'

        tx = x1 + nx*c_off + ax*l/2 + nx*text_offset
        ty = y1 + ny*c_off + ay*l/2 + ny*text_offset
        for text_color, stroke in [('white', 'white'), (text_color, 'none')]:
            yield Tag('text', children=[text_format.format(l)],
                        x=tx, y=ty,
                        font_family='sans-serif', font_size=f'{font_size:.3f}px',
                        text_anchor='middle', dominant_baseline=baseline,
                        transform=f'rotate({angle},{tx},{ty})',
                        stroke=stroke, stroke_width=1,
                        stroke_linejoin='round',
                        fill=text_color)

    crosshairs = lambda pt, s=2, color='black', xf='': Tag('path', fill='none', stroke_width=0.2, stroke=color, transform=xf,
                    d=f'M {pt[0]-s},{pt[1]} h {2*s} m {-s},{s} v {-2*s}')

    current_point = (0, 0)
    next_point = (offset, 0)

    tags.append(crosshairs(current_point, color=current_color(40, 30)))
    tags.append(crosshairs(next_point, color=next_color(40, 30), xf=f'rotate(45,{next_point[0]},{next_point[1]})'))

    if start == 'offset':
        current_point, next_point = next_point, current_point
        current_color, next_color = next_color, current_color
        current_fill, next_fill = next_fill, current_fill

    dim_tags += dimension(current_point, next_point, h=5, offset_1=2)
    dim_tags += dimension(current_point, (current_point[0]-initial_radius, current_point[1]), h=5, offset_1 = 2)
    dim_tags += dimension((current_point[0]+initial_radius, current_point[1]), next_point, h=10, offset_1=mesh_width/2)

    #if (r := offset - mesh_space - initial_radius) > 0:
    #    ring(next_point, next_color)

    radii = []
    for i in reversed(range(steps)):
        print(f'Radius at step {i}: {current_radius:.2f} mm')
        radii.append(current_radius)
        old_radius = current_radius
        ring(current_point, current_color, current_fill)

        sign = -1 if current_point[0] - next_point[0] < 0 else 1
        px, py = current_point
        dim_tags += dimension((px + sign*current_radius, 0), (px + sign*old_radius, 0), h=10, offset_2=mesh_width/2)
        if i > 0:
            dim_tags += dimension((px + sign*current_radius, 0), (px + sign*(current_radius+mesh_space), 0), h=10 + i*3,
                                  offset_2=mesh_width/2)

        if i < 2:
            dim_tags += dimension((px + sign*current_radius, 0), (px - sign*current_radius, 0), h=current_radius + 5,
                                  offset_2=mesh_width/2)

        old_radius = current_radius
        current_radius += mesh_space + offset
        current_point, next_point = next_point, current_point
        current_color, next_color = next_color, current_color
        current_fill, next_fill = next_fill, current_fill

        px, py = current_point
        if i > 0:
            dim_tags += dimension((px - sign*(old_radius-offset), 0), (px - sign*(current_radius-2*offset), 0), h=10)
            dim_tags += dimension((px - sign*current_radius, 0), (px - sign*(current_radius-2*offset), 0), h=10,
                                  offset_1=mesh_width/2)


    radii.append(radii[-1] + 20)
    h = 20 - wire_diameter
    h_off = -h/2 - wire_diameter/2
    points = [(0, 0, +h_off)]
    for i, r in enumerate(radii):
        sign_i = (i%2)*2 - 1
        offset_i = 0 if i%2 == 0 else offset
        a = math.pi/3 if i == 0 else -sign_i * math.pi/2

        x = offset_i + math.cos(a)*r
        y =            math.sin(a)*r
        z = sign_i*h/2 + h_off
        points.append((x, y, z))

        a = -sign_i * math.pi/2
        a -= math.pi/4
        r += 15
        x = offset_i + math.cos(a)*r
        y =            math.sin(a)*r
        z = sign_i*h/2 + h_off
        points.append((x, y, z))

        if i == len(radii)-1:
            break

        a -= math.pi/2
        x = offset_i + math.cos(a)*r
        y =            math.sin(a)*r
        z = -sign_i*h/2 + h_off
        points.append((x, y, z))

    lens_cum = np.zeros(len(points))
    total = 0
    for i, (p1, p2) in enumerate(zip(points[:-1], points[1:])):
        total += math.dist(p1, p2)
        lens_cum[i+1] = total
    lens_cum /= lens_cum[-1]

    points = scipy.interpolate.make_interp_spline(lens_cum, points)(np.linspace(0, 1, 500))

    if out_stl:
        with tempfile.NamedTemporaryFile('w', suffix='.scad') as f:
            points_array = ', '.join(f'[{x:.3f},{y:.3f},{z:.3f}]' for x, y, z in points)
            f.write(f'''
                points = [{points_array}];

                $fn = 17;
                d = {wire_diameter};

                rotate([0, 0, 270])
                for (i=[0:500-2])
                    hull() {{
                        translate(points[i]) sphere(d=d);
                        translate(points[i+1]) sphere(d=d);
                    }}
            ''')
            f.flush()
            print(f'wrote {len(points)} points, rendering')
            openscad = os.environ.get('OPENSCAD', 'openscad')
            openscad_flags = os.environ.get('OPENSCAD_FLAGS', '').split()
            subprocess.run([openscad, '-o', str(out_stl), *openscad_flags, f.name], check=True)

    tags.append(Tag('path', fill='none', stroke='magenta', stroke_width='0.5mm', stroke_linecap='round', stroke_linejoin='round',
                    d='M ' + ' L '.join(f'{x:.3f},{y:.3f}' for x, y, z in points)))

    r_max = current_radius + offset
    out_svg.write_text(str(Tag.setup_svg(tags + dim_tags, bounds=((-r_max, -r_max), (+r_max, +r_max)), margin=3)))


if __name__ == '__main__':
    cli()