import numpy as np
import pigpio, time

class Gcode(object):

    def __init__(self, text):
        self.data = []
        self.count = 0
        self.lines = text.split('\n')
        self.start_lines, self.print_lines, self.end_lines = self.split_lines()

    def split_lines(self):
        def m107count(lines):
            count = 0
            for i, line in enumerate(lines):
                key = line.split(' ')[0]
                if key == 'M107':
                    count = count + 1
                if count > 1:
                    return i
        
        start = m107count(self.lines) + 1
        end = len(self.lines) - m107count(self.lines[::-1]) - 1

        start_lines = self.lines[:start]
        print_lines = self.lines_to_ndarray(self.lines[start:end])
        end_lines = self.lines[end:]

        return start_lines, print_lines, end_lines

    def lines_to_ndarray(self, print_lines):
        arr = [[0, 0, 0, 0, 0, 0]]
        keys = []

        for line in print_lines:
            splited = line.split(' ')
            d = { s[0].strip() : float(s[1:]) for s in splited[1:] }
            a = []
            for i, key in enumerate(['X', 'Y', 'Z', 'E', 'F', 'S']):
                if key in d:
                    a.append( d.get(key) )
                else:
                    a.append( arr[-1][i] )
            keys.append( splited[0] )
            arr.append(a)

        return [ keys, np.array(arr[1:]) ]

def mm_to_pulse_factor_for_belt_pulley(belt_pitch, pulley_tooth, motor_pulse_rate, microstep=1):
    '''
    belt_pitch [mm]
    pulley_tooth [T]
    motor_pulse_rate [pulse/rev]
    return []
    '''
    feed = belt_pitch * pulley_tooth # mm/rev
    factor = motor_pulse_rate * microstep / feed # [pulse/rev] / [mm/rev] = [pulse/mm]
    return factor

def get_pulse_counts_from_gcode(filepath):
    with open(filepath, 'r') as f:
        gcode = Gcode( f.read() )

    # delete other than xyz
    gcode = np.delete( gcode.print_lines[1], [3, 4, 5], 1 )

    # diff elements
    gcode_diff = np.diff(gcode, n=1, axis=0)
    
    # convert factor mm to pulse count
    array_length = gcode_diff.shape[0]
    factor_array = np.array([
        np.full( array_length, mm_to_pulse_factor_for_belt_pulley(2, 16, 200, 8) ),
        np.full( array_length, mm_to_pulse_factor_for_belt_pulley(2, 16, 200, 8) ),
        np.full( array_length, mm_to_pulse_factor_for_belt_pulley(2, 16, 200, 8) )
    ]).T

    # mm to pulse count
    pulse_array = gcode_diff * factor_array

    return np.round(pulse_array)

def pulse_count_to_pulse_array(pulse_count):
    count = int(max(np.abs(pulse_count)))
    if count == 0:
        return np.array([[0], [0], [0]])
    index = np.arange(count)
    arr = []
    for i in pulse_count:
        if i == 0:
            arr.append( np.zeros(count, np.int8) )
        else:
            tmp = np.mod( index, round(count / abs(i), 0) )
            tmp_max = tmp.max()
            if tmp_max == 0:
                arr.append( np.full(count, 1 * np.sign(i), np.int8) )
            else:
                tmp = tmp / tmp.max() * np.sign(i)
                arr.append( tmp.astype(np.int8) )
    return np.array(arr)

def move_motor(X, X_DIR, Y, Y_DIR, ENABLE, pulse_array, velocity):
    pi = pigpio.pi()
    pi.set_mode(X, pigpio.OUTPUT)
    pi.set_mode(Y, pigpio.OUTPUT)
    pi.set_mode(X_DIR, pigpio.OUTPUT)
    pi.set_mode(Y_DIR, pigpio.OUTPUT)
    pi.set_mode(ENABLE, pigpio.OUTPUT)

    pi.write( ENABLE, 0 )

    try:
        print(pulse_array.shape)
        for pulse in pulse_array:
            x_val, y_val = pulse[0], pulse[1]
            pi.write( X_DIR, int((np.sign(x_val) + 1) / 2) )
            pi.write( Y_DIR, int((np.sign(y_val) + 1) / 2) )
            pi.write( X, abs(x_val) )
            pi.write( Y, abs(y_val) )
            time.sleep(velocity)
            pi.write(X, 0)
            pi.write(Y, 0)
            time.sleep(velocity)
    except(KeyboardInterrupt): # when ctrl+c is pressed
        pi.write(X, 0)
        pi.write(Y, 0)
        pi.write( ENABLE, 1 )
        pi.stop()
        print('\nstop motor')

def main():
    print('calcurate pulse array')
    pulse_counts = get_pulse_counts_from_gcode('programs/data/XYZcube.gco')
    pulse_array = np.concatenate([ pulse_count_to_pulse_array(pulse_count) for pulse_count in pulse_counts], 1)
    pulse_array = pulse_array.T

    print('move motor')
    move_motor(
        X = 27, 
        X_DIR = 17, 
        Y = 23, 
        Y_DIR = 22, 
        ENABLE = 24, 
        pulse_array = pulse_array, 
        velocity = 0.0001
    )

if __name__ == '__main__':
    main()