"""
write_blazed.py
Andrew Lorimer, January 2025
Monash University

Writes a reflective brazed grating pattern (sawtooth) using
the confocal direct laser writing setup on Sb2S3 thin-film PCM.
"""

import configparser as cp
import cv2
import numpy as np
import sys
import matplotlib.pyplot as plt
#import write_path
#from pipython import datarectools, pitools
#import nidaqmx
#import pipython
import time

def calculate_offset_vertices(outer_vertices, offset_distance):
    """
    Calculate the vertices of an inner triangle offset inward from an outer triangle.
    
    Parameters:
    - outer_vertices: A list of tuples [(x1, y1), (x2, y2), (x3, y3)] defining the outer triangle.
    - offset_distance: The distance by which the inner triangle is offset inward.
    
    Returns:
    - A list of tuples [(x1', y1'), (x2', y2'), (x3', y3')] defining the inner triangle.
    """
    def unit_vector(vector):
        """Return the unit vector from two points"""
        return vector / np.linalg.norm(vector)

    def normal(vector):
        """Return the normal vector from two points"""
        return np.array([-vector[1], vector[0]])
    
    def intersection(p1, p2, p3, p4):
        """Find the intersection point of two lines defined by points (p1, p2) and (p3, p4)"""
        A1, B1 = p2[1] - p1[1], p1[0] - p2[0]
        C1 = A1 * p1[0] + B1 * p1[1]
        
        A2, B2 = p4[1] - p3[1], p3[0] - p4[0]
        C2 = A2 * p3[0] + B2 * p3[1]
        
        determinant = A1 * B2 - A2 * B1
        if determinant == 0:
            raise ValueError("Lines do not intersect")
        
        x = (B2 * C1 - B1 * C2) / determinant
        y = (A1 * C2 - A2 * C1) / determinant
        return (x, y)

    def is_parallel(vec1, vec2):
        uv1 = unit_vector(vec1)
        uv2 = unit_vector(vec2)
        return np.all(abs(uv1 - uv2) < 0.001)
        
    inner_vertices = []
    num_vertices = len(outer_vertices)
    
    for i in range(num_vertices):
        # Current vertex and next vertex
        p1 = np.array(outer_vertices[i])
        p2 = np.array(outer_vertices[(i + 1) % num_vertices])
        
        # Calculate edge direction and inward normal
        edge_vector = p2 - p1
        inward_normal = unit_vector(normal(edge_vector)) * offset_distance
        
        # Offset the two points on the edge
        offset_p1 = p1 + inward_normal
        offset_p2 = p2 + inward_normal
        
        # Store the offset line
        inner_vertices.append((offset_p1, offset_p2))
    
    # Find intersections of the offset lines
    result_vertices = []
    for i in range(num_vertices):
        # Current offset line and next offset line
        line1 = inner_vertices[i]
        line2 = inner_vertices[(i + 1) % num_vertices]
        
        # Calculate the intersection of the two lines
        int = intersection(line1[0], line1[1], line2[0], line2[1])
        result_vertices.append(int)

    result_vertices = np.array([result_vertices[2], result_vertices[0], result_vertices[1]])
    
    result_edges = np.array([
        [result_vertices[0], result_vertices[1]],
        [result_vertices[1], result_vertices[2]],
        [result_vertices[2], result_vertices[0]]
    ])
    outer_edges = np.array([
        [outer_vertices[0], outer_vertices[1]],
        [outer_vertices[1], outer_vertices[2]],
        [outer_vertices[2], outer_vertices[0]]
    ])

    if not is_parallel(outer_vertices[1] - outer_vertices[0], result_vertices[1] - result_vertices[0]):
        raise ValueError("Offset is too large for the given outer trianlge")

    return result_vertices

if __name__ == "__main__":

    #task, pidevice = write_path.setup()

    # Input parameters
    grating_height = None # h, nm
    grating_pitch = 2000   # Lambda, nm
    #blaze_angle = None # theta_b, degrees
    n_blazes = 10      # number of steps ("teeth")
    beam_dia = 500     # diameter of laser beam, nm
    base_height = 0  # nm
    wavelength = 450    # nm
    order = 1
    n_cr = 1.3

    blaze_angle = np.arcsin(order*wavelength/(2*grating_pitch)) / np.pi * 180
    print("Calculated blaze angle: {:.1f} deg".format(blaze_angle))

    grating_height = np.tan(blaze_angle*np.pi/180)*grating_pitch
    print("Calculated grating height: {:.1f} nm".format(grating_height))

    diff_angle = np.arcsin(wavelength / grating_pitch) / np.pi * 180
    print("First order diffraction angle: {:.1f} deg".format(diff_angle))

    # Calculate vertices of a single step ("tooth")
    if blaze_angle is None:
        d1 = grating_pitch # axial length of upwards section of sawtooth
        d2 = 0 # axial length of downwards section of sawtooth
        back_angle = 90 # angle between grating normal and downwards section of sawtooth, degrees
    else:
        d1 = grating_height / np.tan(blaze_angle * np.pi / 180) # axial length of upwards section of sawtooth
        d2 = grating_pitch - d1 # axial length of downwards section of sawtooth
        back_angle = np.arctan(d2 / grating_height) / np.pi * 180 # angle between grating normal and downwards section of sawtooth, degrees
    
    # Put these vertices into an array as coordinates [ [x1, y1], [x2, y2], [x3, y3] ]
    pts = [
        [0,             0               ],
        [d1,            grating_height  ],
        [grating_pitch, 0               ]
    ]

    # Calculate vertices of the inner triangles to fill in the middle
    n_passes = 1
    while True:
        try:
            pts_inner = calculate_offset_vertices(np.array(pts[-3:]), -beam_dia)
        except ValueError:
            # Not possible to fit any smaller triangles in the middle with given beam diameter
            break
        # Add inner triangles to the list of vertices
        for pt_inner in pts_inner:
            pts.append(pt_inner)
        n_passes += 1

    lines = []

    # Generate lines between each set of three consecutive vertices
    for i in range(len(pts) // 3):
        lines.append((pts[i*3+0], pts[i*3+1]))
        lines.append((pts[i*3+1], pts[i*3+2]))
        lines.append((pts[i*3+2], pts[i*3+0]))
    lines = np.array(lines)

    # Remove a bit of the line in the corner to prevent overwriting
    lines[2,1,0] += beam_dia


    # Duplicate this set of triangles a number of times along the axis of the grating (defined by n_blazes)
    blazes_addition_x = np.tile(np.repeat(np.arange(n_blazes) * grating_pitch, 3*n_passes), (2, 1)).T
    blazes_addition_y = np.zeros((n_blazes * 3 * n_passes, 2))
    blazes_addition = np.stack((blazes_addition_x, blazes_addition_y), axis=-1)
    lines = np.tile(lines, (n_blazes, 1, 1)) + blazes_addition

    # Calculate a set of y values for lines to make up the solid rectangular base
    base_y_vals = np.arange(-base_height, 0, beam_dia)

    # Calculate alternating x values for a zig-zag raster scan
    base_x_vals_single = np.array([0, grating_pitch*n_blazes])
    base_x_vals_duplicated = np.array((base_x_vals_single, np.flip(base_x_vals_single)))
    base_x_vals = np.tile(base_x_vals_duplicated, (int(len(base_y_vals)//2), 1))
    # Deal with odd number of lines
    if len(base_y_vals) % 2 != 0:
        base_x_vals = np.vstack((base_x_vals, base_x_vals_single))

    # Put x and y values into a list of pairs of points
    base_y_vals = np.tile(base_y_vals, (2, 1)).T
    base_lines = np.stack((base_x_vals, base_y_vals), axis=2)

    # Add these lines to the overall array
    lines = np.concatenate((lines, base_lines))

    # Plot
    fig, ax = plt.subplots()
    ax.set_aspect("equal")
    for pair in lines:
        ax.plot(pair[:,0], pair[:,1], '-', color='black')  
    plt.show()