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

Converts a path in a monochrome image to a series of points 
in a defined format for laser writing with the confocal 
setup. Configuration is done in img_to_path.ini. This version
for use with write_path.py (Python).
"""

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

DEFAULTS = {'dimensions': '100, 100',
            'z': '0',
            'n_levels': '10',
            'pitch': 1,
            'black_exposure': '1'}
DEFAULT_CONFIG_PATH = "img_to_path_2.ini"

if __name__ == "__main__":
    
    # Get config file path from arguments or use default path
    if len(sys.argv) > 1 and sys.argv[1]:
        config_path = sys.argv[1]
    else:
        config_path = DEFAULT_CONFIG_PATH

    # Set up config parser
    config = cp.ConfigParser()
    config['DEFAULT'] = DEFAULTS
    config.read_file(open(config_path))

    # Set config values
    real_dim = np.fromstring(config.get('Main', 'dimensions'), dtype=float, sep=",")
    z = config.getfloat('Main', 'z')
    n_levels = int(config.getfloat('Main', 'n_levels'))
    pitch = config.getfloat('Main', 'pitch')
    black_exposure = config.getfloat('Main', 'black_exposure')
    input_path = config.get('Main', 'input')

    # Read image
    im = cv2.imread(input_path, cv2.IMREAD_GRAYSCALE)
    im_dim = im.shape

    # Check if aspect ratios match
    if (im_dim[0] / im_dim[1] != real_dim[0]/real_dim[1]):
        print("Warning: input and output aspect ratios do not match - result will be distorted")

    # Convert image to defined number of levels
    im = im // (255/n_levels)

    # Upsample/downsample image for the writing resolution
    n_points = real_dim / pitch
    n_points = n_points.astype(int)
    im = cv2.resize(im, (int(n_points[0]), int(n_points[1])), interpolation=cv2.INTER_CUBIC)

    # Expand dynamic range
    im = (im - im.min()) * 255/(im.max()-im.min())
    # Display image
    fig, ax = plt.subplots()
    ax.imshow(im, cmap='gray')
    plt.show()

    print("Min: {}, Max: {}".format(im.min(), im.max()))

    min_exposure = 0.1
    #exposure_times = (100**((n_levels - im)/n_levels) - (1 - min_exposure)) * black_exposure / (100 - (1 - min_exposure))
    p = 1
    #exposure_times = (black_exposure / min_exposure)**(((n_levels - im)/n_levels)**p) * min_exposure
    exposure_times = -(black_exposure - min_exposure)/n_levels * im + black_exposure
    fig, ax = plt.subplots()
    ax.imshow(exposure_times, cmap='gray')
    plt.show()

    print("Min: {}, Max: {}".format(exposure_times.min(), exposure_times.max()))

    cont = input("Write image? (y/n): ")

    if cont.lower() != 'y':
        sys.exit()

    print("Writing image")

    task, pidevice = write_path.setup()

    start_x = 0
    start_y = 3
    pidevice.MOV({1: start_x, 2: start_y, 3: z})
    

    try:
        for i in range(34, n_points[0]):
            for j in range(n_points[1]):
                pidevice.MOV({1: i*pitch+start_x, 2: j*pitch+start_y})
                pitools.waitontarget(pidevice, [1, 2])
                time.sleep(0.1)
                exposure_time = exposure_times[i, j]
                print(exposure_time)
                if exposure_time >= min_exposure:
                    task.write(True)
                    time.sleep(exposure_time)
                    task.write(False)
                    time.sleep(0.1)
    finally:
        print("Exiting")
        pidevice.MOV({1: start_x, 2: start_y, 3: z})
        pitools.waitontarget(pidevice, [1, 2])
        pitools.stopall(pidevice)
        task.write(False)
        task.close()