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

Writes a transmission binary grating pattern 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

TRANSITION_SPEED = 100   # speed for moving between lines, um/s

# Input parameters
grating_pitch = 2.7 # Lambda, um
beam_width = 0.7
writing_speed = 10  # speed for writing crystallised lines, um/s
dimensions = (50, 50)   # horizontal * vertical, um
origin = (0, 0) # horizontal * vertical, um

def wait_target(pidevice, thresh=0.01, axes=[1, 2, 3]):
    target = pidevice.qMOV(axes)
    on_target = False
    while on_target == False:
        pos = pidevice.qPOS(axes)
        on_target = True
        distances = np.abs(np.array(list(pos.values())) - np.array(list(target.values())))
        on_target = np.all(distances < thresh)
            

def main(task, pidevice):

    # One block is a solid area with constant refractive index in the 
    # diffraction grating
    # One line is a single pass of the laser, with many lines making up a 
    # block (depending on the beam width)

    n_blocks = int(dimensions[0] // grating_pitch)
    n_lines = int(grating_pitch // beam_width)

    pidevice.VCO({1: True})

    updown = 0

    for i in range(n_blocks):
        if (i % 2 == 0):
            # Even number block - crystallise

            y_block = origin[0] + i * grating_pitch

            for j in range(n_lines):
                y_line = y_block + j*beam_width
                x_start = origin[1] + int(not updown)*dimensions[1]
                x_end = origin[1] + updown*dimensions[1]

                # Go to line start position
                pidevice.VEL({1: TRANSITION_SPEED})
                pidevice.MOV({1: x_start, 2: y_line})
                wait_target(pidevice)
                time.sleep(0.3)
                
                # Write line
                pidevice.VEL({1: writing_speed})
                pidevice.MOV({1: x_end, 2: y_line})
                #time.sleep(0.05)   # compensation for acceleration period
                task.write(True)
                wait_target(pidevice)
                task.write(False)
                time.sleep(0.3)

                updown = int(not updown)
        else:
            # Odd number line - do nothing
            pass        

if __name__ == "__main__":
    task, pidevice = write_path.setup()
    try:
        main(task, pidevice)
    finally:
        print("Exiting")
        task.write(False)
        task.close()
        pidevice.VCO({1: True, 2: True})
        pidevice.VEL({1: TRANSITION_SPEED, 2: TRANSITION_SPEED})
        pidevice.MOV({1: origin[1], 2: origin[0]})
        pitools.waitontarget(pidevice, [1, 2])
        pitools.stopall(pidevice)