import machine
import utime
import time
import math
import sys
import select
from led import *
from config import *
import _thread
from TMC2130 import TMC2130_Rotate, TMC2130_Init, rotate_motor_test


lens_cs = machine.Pin(4, machine.Pin.IN)
body_cs = machine.Pin(5, machine.Pin.OUT)
vd_lens = machine.Pin(2, machine.Pin.OUT)
tx_pin = machine.Pin(0, machine.Pin.OUT)



LED_OK.value(1)

poll_obj = ''

focal_length = 0
aperture_raw = 0
focus = 0
focus_limit = 0
status_sent = False

counter = 0 #VD_LENS pulse counter
ap = 4884 #aperture
ap_odd = False #variable for controlling one bit in send messages

focus_pos = 0 #focus position to be set manually

af_in_progress = False
focus_mode = "MF" #AF/MF

focus_pos = 0
change_focus = False
relative_focus = False

utime.sleep(1)
vd_lens.value(1)
body_cs.value(0)
tx_pin.value(0)

utime.sleep(0.1)

vd_lens.value(0)
tx_pin.value(1)
time.sleep_ms(6)
tx_pin.value(0)
time.sleep_ms(106)
tx_pin.value(1)
time.sleep_ms(64)
tx_pin.value(0)
time.sleep_ms(80)
vd_lens.value(1)
time.sleep_ms(20)

rx_ack = False

def Zoom(focal_distance, steps_per_iter):
    dir = 0
    reverse = False
    
    if focal_distance > focal_length:
        dir = 1
        while focal_distance > focal_length:
            TMC2130_Rotate(dir, steps_per_iter, ZOOM_SPEED)
    else:
        while focal_distance < focal_length:
            TMC2130_Rotate(dir, steps_per_iter, ZOOM_SPEED)

def Comms_Init():
    global poll_obj
    
    poll_obj = select.poll()
    poll_obj.register(sys.stdin, select.POLLIN)
    
def Comms_Read():
    global poll_obj, ap, focus_mode, change_focus, focus_pos, relative_focus
    
    poll_results = poll_obj.poll(1)
    if poll_results:
        rw = sys.stdin.readline().strip()
        rw = rw.upper()
        if "APERTURE" in rw:
            try:
                data = rw.split(' ')
                if "REL" in rw:
                    ap = aperture_raw + int(data[1])
                else:                
                    ap = int(data[1])
            except:
                print("APERTURE SET ERROR")
        if "FOCUS" in rw:
            try:
                data = rw.split(' ')
                focus_pos = int(data[1])
                change_focus = True
                if "REL" in rw:
                    relative_focus = True
            except:
                print("FOCUS SET ERROR")
        if "AF" in rw:
            focus_mode = "AF"
        if "MF" in rw:
            focus_mode = "MF"
        if "ZOOM" in rw:
            try:
                data = rw.split(' ')
                if "STEPS" in rw:
                    if int(data[1]) < 2:
                        if int(data[1]) >= 0:
                            if int(data[2]) > 0:
                                _thread.start_new_thread(TMC2130_Rotate, (int(data[1]), int(data[2]), ZOOM_SPEED))
                else:
                    if int(data[1]) > FOCAL_LEN_MIN:
                        if int(data[1]) < FOCAL_LEN_MAX:
                            if int(data[2]) > 0:
                                _thread.start_new_thread(Zoom, (int(data[1]), int(data[2])))
            except:
                print("ZOOM SET ERROR")

def focal_length_conv(raw):
    return (math.exp((raw+92063.3)/8082.9)+406.979)/3256.82
    
def aperture_count(value, leftMin=2.8, leftMax=16, rightMin=4884, rightMax=6132):
    leftSpan = leftMax - leftMin
    rightSpan = rightMax - rightMin

    valueScaled = float(value - leftMin) / float(leftSpan)
    return rightMin + (valueScaled * rightSpan)

def aperture(value):
    if value < AP_MIN:
        value = AP_MIN
    if value > AP_MAX:
        value = AP_MAX
    return value
    
def rx_end(pin):
    rx_ack = True
    
def vd_pulse():
    global counter
    vd_lens.value(0)
    time.sleep_us(30)
    vd_lens.value(1)
    if counter > 239:
        counter = 1
    else:
        counter += 1
    
def print_rx():
    global focal_length, aperture_raw, focus, focus_limit, counter
    
    msg = b''
    msg_len = 0
    index = 0
    read_in_progress = False
    
    if uart_bodytx.any():
        received_data = uart_bodytx.read()
        for b in received_data:
            if read_in_progress:
                msg += b.to_bytes(1, 'little')
                if index == 2:
                    msg_len = b
                if index == 3:
                    msg_len = msg_len + (b << 8)
                if index == msg_len:
                    read_in_progress = False
                    if msg[5] == 6:
                        focal_length = focal_length_conv(msg[17] + (msg[18] << 8))
                        focus = msg[8] + (msg[9] << 8)
                        if msg[6] == 0x02:
                            focus_limit = "NO"
                        elif msg[6] == 0x0A:
                            focus_limit = "NEAR"
                        elif msg[6] == 0x12:
                            focus_limit = "FAR"
                        else:
                            focus_limit = "UNKNOWN"
                    if msg[5] == 5:
                        aperture_raw = msg[6] + (msg[7] << 8)
                    print("Focus mode:", focus_mode, "Focus limit:", focus_limit, "Focus:", focus, "Focal length:", focal_length, "Aperture:", aperture_raw)
                    LED_MSG.toggle()
                
            else:
                if b == 0xf0:
                    index = 1
                    read_in_progress = True
                    msg = b''
                    msg += b.to_bytes(1, 'little')
            index += 1
        
def tx(data):
    body_cs.value(1)
    uart_bodytx.write(data)
    while not uart_bodytx.txdone():
        pass
    body_cs.value(0)
    time.sleep_us(800)

def send_msg(type, payload):
    global counter
    start = b'\xf0'
    end = b'\x55'
    length = len(payload)+6+3
    checksum = 0
    normal_msg = b'\x01'
    msg = b''
    
    msg = msg+length.to_bytes(2, 'little')+normal_msg+counter.to_bytes(1, 'little')+type+payload
    
    for b in msg:
        checksum += b
        
    msg += checksum.to_bytes(2, 'little')
    
    tx(start+msg+end)
    
def send_aperture(ap):
    global ap_odd
    
    aper = int(aperture(ap))
    
    payload = b'\x3f\x75\x00'+aper.to_bytes(2,'little')+aper.to_bytes(2,'little')+b'\x1c\x00\x00\x02'
    if ap_odd:
        ap_odd = False
        payload += b'\x00\x00'
    else:
        ap_odd = True
        payload += b'\x01\x01'
    payload += b'\x02\x00\x02\x01\x00\x00\x00\x2f\x15\x17'
    send_msg(b'\x03', payload)

def send_af():
    global ap, ap_closing, af_in_progress, focus_mode
    
    vd_pulse()
    time.sleep_ms(10)
    send_aperture(ap)
    #jump out and move in
    #send_msg(b'\x04', b'\x00\x00\x19\x83\x00\x00\x66\x00\x00\x80\x41\x00\x00\x2f\x09\x0a\x1f\x06\x02\x83\x88\xd5\x4b\xa4\x3f\x7f\x53\x03\x00\x00')
    #jump in and move out
    #send_msg(b'\x04', b'\x00\x00\x19\x83\x00\x00\x66\x00\x00\x80\x41\x00\x00\x2f\x09\x0a\x1f\x02\x02\x83\x88\xd5\x4b\xa4\x3f\x7f\x53\x03\x00\x00')
    send_msg(b'\x04', b'\x00\x00\x19\x83\x00\x00\x66\x00\x00\x00\x41\x00\x00\x2f\x09\x0a\x1f\x02\x02\x83\x84\xd5\x54\xf2\x3e\x1f\x53\x03\x00\x00')
    af_in_progress = True
    time.sleep_ms(6)
    vd_pulse()
    time.sleep_ms(10)
    send_aperture(ap)
    send_focus(focus_mode, af_in_progress)
    time.sleep_ms(6)
    
def move_focus(position):
    global ap, ap_closing, af_in_progress, focus_mode, focus_pos, FOCUS_MIN, FOCUS_MAX
    
    if position < FOCUS_MIN:
        position = FOCUS_MIN
    if position > FOCUS_MAX:
        position = FOCUS_MAX
    
    vd_pulse()
    time.sleep_ms(10)
    send_aperture(ap)
    send_msg(b'\x04', b'\x00\x00\x19\x83\x00\x00\x66\x00\x00\x00\x41\x00\x00\x2f\x00\x00\x1d'+position.to_bytes(2, 'little')+b'\x00\x40')
    af_in_progress = False
    time.sleep_ms(6)
    vd_pulse()
    time.sleep_ms(10)
    send_aperture(ap)
    send_focus(focus_mode, af_in_progress)
    time.sleep_ms(6)
    
def send_focus(mode, af_in_progress):
    payload = b''
    mode_sw = b'\x81'
    if mode == "AF":
        mode_sw = b'\x83'
    
    if af_in_progress:
        payload = b'\x00\x00\x19'+mode_sw+b'\x00\x00\x66\x00\x00\x80\x41\x00\x00'
    else:
        payload = b'\x00\x00\x19'+mode_sw+b'\x00\x00\x66\x00\x00\x00\x01\x00\x00'
        
    send_msg(b'\x04', payload)

def shutdown():
    vd_pulse()
    time.sleep_ms(10)
    print_rx() #00 00 198300 00 f600 00 00 09 00 00 1d 0000 002cc80255
    tx(b'\xf0\x1b\x00\x01\xc4\x04\x00\x00\x19\x83\x00\x00\xf6\x00\x00\x00\x09\x00\x00\x1d\x00\x00\x00\x2c\xc8\x02\x55')
    tx(b'\xf0\x16\x00\x01\x1c\x04\x00\x00\x19\x81\x00\x00\x66\x00\x00\x00\x01\x00\x00\x38\x01\x55') #Focus??
    time.sleep_ms(6)

def normal_loop():
    global ap, ap_closing, af_in_progress, focus_mode, focus_pos, change_focus, relative_focus
    
    if change_focus:
        if relative_focus:
            move_focus(focus+focus_pos)
            relative_focus = False
        else:
            move_focus(focus_pos)
        change_focus = False
    else:    
        vd_pulse()
        time.sleep_ms(10)
        send_aperture(ap)
        send_focus(focus_mode, af_in_progress)
        time.sleep_ms(6)
        print_rx()
        vd_pulse()
        time.sleep_ms(10)
        send_aperture(ap)
        send_focus(focus_mode, af_in_progress)
        time.sleep_ms(6)
        print_rx()
    

TMC2130_Init()
Comms_Init()
lens_cs.irq(handler=rx_end, trigger=machine.Pin.IRQ_FALLING)

vd_pulse()
body_cs.value(1)
time.sleep_ms(4)
uart_bodytx = machine.UART(0, baudrate=750000)  # UART1, GP0 (TX) and GP1 (RX) #, tx=machine.Pin(16), rx=machine.Pin(17)
body_cs.value(0)
time.sleep_ms(15)
vd_pulse()
time.sleep_ms(20)
vd_pulse()

uart_bodytx.read()

tx(b'\xf0\x0b\x00\x02\x00\x0b\x60\x00\x78\x00\x55')
print_rx()
tx(b'\xf0\x0d\x00\x02\x00\x09\x00\x00\x00\x00\x18\x00\x55')
print_rx()
tx(b'\xf0\x0a\x00\x02\x00\x0d\x01\x1a\x00\x55')
print_rx()
tx(b'\xf0\x0a\x00\x02\x00\x10\x1f\x3b\x00\x55') #Self calibration??
print_rx()
time.sleep_ms(10)
for i in range(0,22):
    vd_pulse()
    time.sleep_ms(20)
time.sleep_ms(6)
print_rx()
tx(b'\xf0\x19\x00\x02\x00\x0a\xff\xff\x03\x00\x00\x00\x00\x00\x3f\x00\x00\x00\x00\x00\x00\x00\x65\x02\x55')
print_rx()
time.sleep_ms(10)


while True:
    normal_loop()
    Comms_Read()
    



