"""
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 main(task, pidevice):

    # Input parameters
    grating_pitch = 5000 # Lambda, nm
    n = 1      # number of steps ("teeth")
    beam_dia = 1350 # diameter of laser beam, nm
    wavelength = 450    # nm
    order = 1
    n_cr = 1.3
    length = 12

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

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

    #y_single_step = np.arange(0, grating_pitch, beam_dia)
    y_single_step = []
    speeds_single_step = []
    #for y in y_single_step:
    dia = beam_dia
    y = 0
    while y <= grating_pitch:
        speeds_single_step.append(np.exp(7*y/grating_pitch))
        y_single_step.append(y)
        y += dia
        dia *= 0.9

    print(y_single_step)
    print(speeds_single_step)

    #fig, ax = plt.subplots()
    #ax.plot(y_single_step, speeds_single_step, marker='.', linestyle='None')
    #ax.set_yscale('log')
    #plt.show()

    origin = (0, 0)
    pidevice.VCO({1: True, 2: True})
    pidevice.VEL({1: 100, 2: 100})
    y_step = 0

    for i in range(n):
        pidevice.MOV({1: origin[0], 2: origin[1]})
        pitools.waitontarget(pidevice, [1, 2, 3])

        updown = 0

        pidevice.MOV({1: 10, 2: 0.6})

        for j in range(len(speeds_single_step)):
            pidevice.MOV({1: origin[0]+updown*length, 2: origin[1]+y_step+y_single_step[j]/1000})
            pitools.waitontarget(pidevice, [1, 2, 3])
            time.sleep(0.3)

            pidevice.VCO({1: True, 2: True})
            print(speeds_single_step[j])
            pidevice.VEL({1: speeds_single_step[j], 2: speeds_single_step[j]})
            
            task.write(True)
            pidevice.MOV({1: origin[0]+(not updown)*length, 2: origin[1]+y_step+y_single_step[j]/1000})
            pitools.waitontarget(pidevice, [1, 2, 3])
            task.write(False)
            
            time.sleep(0.3)

            pidevice.VCO({1: True, 2: True})
            pidevice.VEL({1: 100, 2: 100})
            
            updown = int(not updown)

        y_step += grating_pitch/1000


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

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: 100, 2: 100})
        pidevice.MOV({1: 0, 2: 0})
        pitools.waitontarget(pidevice, [1, 2])
        pitools.stopall(pidevice)