#!/usr/bin/env python3
import numpy as np
import matplotlib.pyplot as plt
from dataclasses import dataclass, field
from typing import List
from datetime import datetime
import os
from copy import deepcopy

@dataclass
class Vertex:
    x: float
    y: float
    z: float

@dataclass
class Line:
    vertices: List[Vertex] = field(default_factory=list)
    def __iter__(self):
        return iter(self.vertices)
    def __len__(self):
        return len(self.vertices)
    def append(self, value):
        self.vertices.append(value)

@dataclass
class Multiline:
    lines: List[Line] = field(default_factory=list)
    def __iter__(self):
        return iter(self.lines)
    def append(self, value):
        self.lines.append(value)
    def extend(self, value):
        self.lines.extend(value)

@dataclass
class Stage:
    multiline: Multiline
    name: str = None
    filename: str = None
    lp: float = None
    v: float = None
    xo: float = None
    yo: float = None

@dataclass
class Job:
    stages: List[Stage] = field(default_factory=list)
    def __iter__(self):
        return iter(self.stages)
    def append(self, value):
        self.stages.append(value)

job_file_header = '''
InvertZAxis 1

% Writing configuration
GalvoScanMode
ContinuousMode
PiezoSettlingTime 10
GalvoAcceleration 1
StageVelocity 200

% Scan field offsets
XOffset 0
YOffset 0
ZOffset 0

% Writing parameters
PowerScaling 1.0
ScanSpeed 10000

TextFontSize 10
TimeStampOn
DebugModeOn
MessageOut "starting..."


'''

def draw_job(job: Job):
    for stage in job.stages:
        draw_stage(stage)

def draw_stage(stage: Stage):
    for line in stage.multiline:
        X = []
        Y = []
        Z = []
        for v in line:
            X.append(v.x + stage.xo)
            Y.append(v.y + stage.yo)
            # Z.append(v.z + stage.zo)
        plt.plot(X, Y, markevery=[0,-1], marker='o')


def write_job(filename, job: Job, interface_z:float=0.0):
    with open(filename, "w") as fd:
        fd.write(job_file_header)
        x = 0
        y = 0
        for stage in job:
            dx = stage.xo - x
            dy = stage.yo - y

            fd.write(f"MoveStageX {dx:.3f}\n")
            fd.write(f"MoveStageY {dy:.3f}\n")

            if stage.name != None:
                fd.write(f"MoveStageY -130\n")
                fd.write(f'WriteText "{stage.name}"\n')
                fd.write(f"MoveStageY 130\n")

            if stage.lp != None:
                fd.write(f"LaserPower {stage.lp:.1f}\n")

            fd.write(f"FindInterfaceAt {interface_z:.3f}\n")
            print(f"movestage {dx:.3f} {dy:.3f}")
            x = stage.xo
            y = stage.yo

            stage_filename = write_stage(stage)
            fd.write(f"include {stage_filename}\n")

        dx = -x
        dy = -y
        fd.write(f"MoveStageX {dx:.3f}\n")
        fd.write(f"MoveStageY {dy:.3f}\n")

def write_stage(stage: Stage) -> str:
    if not stage.filename.endswith('.gwl'):
        stage.filename += '.gwl'
    if not os.path.exists(stage.filename):
        with open(stage.filename, "w") as fd:
            for line in stage.multiline:
                for v in line:
                    fd.write(f"{v.x:.3f} {v.y:.3f} {v.z:.3f}\n")
                fd.write("write\n")
    return stage.filename

def lines_0_dir(w, p, xo=0, yo=0, zo=0, reverse_direction=False, alternating=False) -> Multiline:
    n = np.arange(0, w, p)
    ml = Multiline()

    direction = True
    if reverse_direction:
        direction = not direction

    for i in n:
        line = [Vertex(xo, yo +i, zo), Vertex(xo+w, yo+i, zo)]
        if not direction:
            line.reverse()
        if alternating:
            direction = not direction
        ml.append(line)
    return ml

def lines_90_dir(w, p, xo=0, yo=0, zo=0, reverse_direction=False, alternating=False) -> Multiline:
    n = np.arange(0, w, p)
    ml = Multiline()

    direction = True
    if reverse_direction:
        direction = not direction

    for i in n:
        line = [Vertex(xo+i, yo, zo), Vertex(xo+i, yo+w, zo)]
        if not direction:
            line.reverse()
        if alternating:
            direction = not direction
        ml.append(line)
    return ml

def lines_0(w, p, xo=0, yo=0, zo=0) -> Multiline:
    n = np.arange(0, w, p)
    ml = Multiline()
    for i in n:
        line = [Vertex(xo, yo +i, zo), Vertex(xo+w, yo+i, zo)]
        ml.append(line)
    return ml

def lines_90(w, p, xo=0, yo=0, zo=0) -> Multiline:
    n = np.arange(0, w, p)
    ml = Multiline()
    for i in n:
        line = [Vertex(xo+i, yo, zo), Vertex(xo+i, yo+w, zo)]
        ml.append(line)
    return ml

def lines_45(w, p, xo=0, yo=0, zo=0) -> Multiline:
    n = np.arange(0, w, p)
    ml = Multiline()
    for i in n:
        l = Line([Vertex(xo, yo+i, zo), Vertex(xo+w-i, yo+w, zo)])
        ml.append(l)
        l = Line([Vertex(xo+i, yo, zo), Vertex(xo+w, yo+w-i, zo)])
        ml.append(l)
    return ml

def lines_135(w, p, xo=0, yo=0, zo=0) -> Multiline:
    n = np.arange(0, w, p)
    ml = Multiline()
    for i in n:
        l = [Vertex(xo+i, yo+w, zo), Vertex(xo+w, yo+i, zo)]
        ml.append(l)
        l = [Vertex(xo, yo+w-i, zo), Vertex(xo+w-i, yo, zo)]
        ml.append(l)
    return ml

def board2x2(w, p, xo=0, yo=0, zo=0) -> Multiline:

    ml = Multiline()
    ml.extend(lines_0(w, p, xo-w, yo-w, zo))
    ml.extend(lines_90(w, p, xo, yo-w, zo))
    ml.extend(lines_90(w, p, xo-w, yo, zo))
    ml.extend(lines_0(w, p, xo, yo, zo))
    return ml

def board2x2_45(w, p, xo=0, yo=0, zo=0) -> Multiline:
    ml = Multiline()
    ml.extend(lines_45(w, p, xo-w, yo-w, zo))
    ml.extend(lines_135(w, p, xo, yo-w, zo))
    ml.extend(lines_135(w, p, xo-w, yo, zo))
    ml.extend(lines_45(w, p, xo, yo, zo))
    return ml

def board2x2_135(w, p, xo=0, yo=0, zo=0) -> Multiline:
    ml = Multiline()
    ml.extend(lines_135(w, p, xo-w, yo-w, zo))
    ml.extend(lines_45(w, p, xo, yo-w, zo))
    ml.extend(lines_45(w, p, xo-w, yo, zo))
    ml.extend(lines_135(w, p, xo, yo, zo))
    return ml

def board_circ(w, p, xo=0, yo=0, zo=0) -> Multiline:
    R = np.arange(p/2, w, p)
    ml = Multiline()
    for r in R:
        line = []
        # seg_len = 2 * np.pi * r
        for phi in np.linspace(0, 2*np.pi, 200):
            v = Vertex(r * np.cos(phi + 42*r*r), r * np.sin(phi + 42*r*r), zo)
            line.append(v)
        ml.append(line)
    return ml

def board_variable_pitch(w, p1, p2, xo=0, yo=0, zo=0) -> Multiline:
    avg = (p1 + p2) / 2
    n = int(w / avg)
    delta = p2 - p1

    xo -= w/2
    yo -= w/2

    x = xo
    ml = Multiline()
    for i in range(n):
        line = [Vertex(x, yo, zo), Vertex(x, yo+w, zo)]
        x += p1 + delta * i/n
        ml.append(line)
    return ml

def streamlines_to_multiline(streamlines, xo=0, yo=0, zo=0) ->  Multiline:
    ml = Multiline()
    for line in streamlines:
        l = Line()
        for x, y in line:
            l.append(Vertex(x + xo, y + yo, zo))
        ml.append(l)
    return ml

def multiline_mirror_x(ml: Multiline) -> Multiline:
    ml2 = Multiline()
    for l in ml:
        l2 = Line()
        for v in l:
            v.x = - v.x
            l2.append(v)
        ml2.append(l2)
    return ml2

def multiline_cut_vertex(ml: Multiline, f) -> Multiline:
    ml_new = Multiline()
    for line in ml:
        l = Line()
        for v in line:
            if f(v):
                l.append(v)
            else:
                if len(l) > 0:
                    ml_new.append(l)
                    l = Line()

        if len(l) > 0:
            ml_new.append(l)
    return ml_new