from machine import Pin, PWM, I2C
import time

# -------------------- Constants / Pins --------------------

# Stepper pins (shared for all harnesses)
STEP_PIN = 3  # Pico GP3 → STEP on all A4988
DIR_PIN  = 2  # Pico GP2 → DIR on all A4988
step_pin = Pin(STEP_PIN, Pin.OUT)
dir_pin = Pin(DIR_PIN, Pin.OUT)
step_pin.value(0)
dir_pin.value(0)

# Buttons
EXECUTE_PIN = 13  # Pico GP13 → execute
PREV_ROW_PIN = 11 # Pico GP11 → previous row
NEXT_ROW_PIN = 4  # Pico GP4 → next row
execute_pin = Pin(EXECUTE_PIN, Pin.IN, Pin.PULL_UP)
prev_row_pin = Pin(PREV_ROW_PIN, Pin.IN, Pin.PULL_UP)
next_row_pin = Pin(NEXT_ROW_PIN, Pin.IN, Pin.PULL_UP)

# Number of harnesses
NUM_HARNESSES = 16

# Steps per harness movement
MOVEMENT_AMOUNT = 250

# RGB LED pins (common-anode)
RGB_RED_PIN   = 12  # Pico GP12
RGB_GREEN_PIN = 9   # Pico GP9
RGB_BLUE_PIN  = 10  # Pico GP10

#I2C display pins
i2c = I2C(0, sda=Pin(0), scl=Pin(1), freq=400000)

# lcd_api.py and i2c_lcd.py must be in your Pico filesystem
from lcd_api import LcdApi
from i2c_lcd import I2cLcd

# I2C address of LCD (commonly 0x27)
LCD_ADDR = 0x27  
lcd = I2cLcd(i2c, LCD_ADDR, 4, 20)  # 4 rows, 20 columns

# -------------------- 74HC595 Shift Register Pins --------------------
DATA_PIN  = 6  # Pico GP6 → DS
CLOCK_PIN = 7  # Pico GP7 → SH_CP
LATCH_PIN = 8  # Pico GP8 → ST_CP

data_pin = Pin(DATA_PIN, Pin.OUT)
clock_pin = Pin(CLOCK_PIN, Pin.OUT)
latch_pin = Pin(LATCH_PIN, Pin.OUT)

# 16-bit shift register state for 2x 74HC595
shift_state = 0xFFFF  # all HIGH = all disabled (active LOW)

# -------------------- Shift Register Functions --------------------

def update_shift_register():
    """Send 16-bit shift_state to 2 daisy-chained 74HC595s."""
    global shift_state
    latch_pin.value(0)  # begin
    for i in range(15, -1, -1):  # MSB first
        bit_val = (shift_state >> i) & 1
        data_pin.value(bit_val)
        clock_pin.value(1)
        clock_pin.value(0)
    latch_pin.value(1)  # latch outputs

def enable_harness(index):
    """Active LOW: set harness output LOW to enable stepper."""
    global shift_state
    shift_state &= ~(1 << index)
    update_shift_register()

def disable_harness(index):
    """Active LOW: set harness output HIGH to disable stepper."""
    global shift_state
    shift_state |= (1 << index)
    update_shift_register()

def enable_row_harnesses(harness_list):
    """Enable only harnesses in list; disable others."""
    global shift_state
    shift_state = 0xFFFF  # disable all first
    for h in harness_list:
        shift_state &= ~(1 << h)
    update_shift_register()

def disable_all_harnesses():
    global shift_state
    shift_state = 0xFFFF
    update_shift_register()

# -------------------- Harness Class --------------------

class Harness:
    def __init__(self, index):
        self.index = index
        self.isRaised = False
        self.isEnabled = False

    def enable(self):
        if not self.isEnabled:
            enable_harness(self.index)
            self.isEnabled = True

    def disable(self):
        if self.isEnabled:
            disable_harness(self.index)
            self.isEnabled = False

# -------------------- Treadling Class --------------------

class Treadling:
    def __init__(self):
        self.rows = []  # list of lists: [[0,2,5], [1,3], ...]
        self.row_colors = []  # RGB tuples
        self.row_color_indices = [] # WIF color index per row
        self.num_rows = 0
        self.num_harnesses = NUM_HARNESSES

    def import_wif(self, filename):
        """Load WIF file and parse treadling, tieup, and colors."""
        treadling_map = {}
        tieup_map = {}
        color_table = {}
        row_color_map = {}
        weft_default_color = None
        section = None

        try:
            with open(filename, "r") as f:
                for raw in f:
                    line = raw.strip()
                    if not line or line.startswith(";"):
                        continue
                    if line.startswith("[") and line.endswith("]"):
                        section = line.upper()
                        continue
                    if section == "[TREADLING]" and "=" in line:
                        r, t = line.split("=")
                        treadling_map[int(r)-1] = int(t)-1
                    elif section == "[TIEUP]" and "=" in line:
                        t, h_list = line.split("=")
                        tieup_map[int(t)-1] = [int(x)-1 for x in h_list.split(",")]
                    elif section == "[COLOR TABLE]" and "=" in line:
                        i, rgb = line.split("=")
                        color_table[int(i)-1] = tuple(int(v) for v in rgb.split(","))
                    elif section in ("[TREADLING COLORS]", "[WEFT COLORS]") and "=" in line:
                        r, c = line.split("=")
                        row_color_map[int(r)-1] = int(c)-1
                    elif section == "[WEFT]" and line.startswith("Color="):
                        weft_default_color = int(line.split("=")[1]) - 1

        except OSError as e:
            print("Failed to open WIF file:", e)
            return False

        max_row = max(treadling_map.keys()) if treadling_map else -1
        self.rows = []
        self.row_colors = []
        self.row_color_indices = []

        for row in range(max_row + 1):
            treadle = treadling_map.get(row)
            harness_list = tieup_map.get(treadle, []) if treadle is not None else []
            self.rows.append(harness_list)

            color_idx = row_color_map.get(row, weft_default_color)
            self.row_color_indices.append(color_idx)
            self.row_colors.append(color_table.get(color_idx, (0,0,0)))  # fallback to black

        self.num_rows = len(self.rows)
        print("WIF import complete:", self.num_rows, "rows")
        return True

    def get_row(self, row):
        if 0 <= row < self.num_rows:
            return self.rows[row]
        return []

    def get_row_color(self, row):
        if 0 <= row < self.num_rows:
            return self.row_colors[row]
        return (0,0,0)

    def get_row_color_index(self, row):
        if 0 <= row < self.num_rows:
            return self.row_color_indices[row]
        return None

# -------------------- RGB LED --------------------

class RGBLed:
    def __init__(self, r_pin, g_pin, b_pin, freq=1000):
        self.red   = PWM(Pin(r_pin), freq=freq, duty_u16=65535)
        self.green = PWM(Pin(g_pin), freq=freq, duty_u16=65535)
        self.blue  = PWM(Pin(b_pin), freq=freq, duty_u16=65535)

    def set_color(self, r, g, b):
        r = max(0, min(255, r))
        g = max(0, min(255, g))
        b = max(0, min(255, b))
        self.red.duty_u16(65535 - (r * 65535 // 255))
        self.green.duty_u16(65535 - (g * 65535 // 255))
        self.blue.duty_u16(65535 - (b * 65535 // 255))

    def off(self):
        self.red.duty_u16(65535)
        self.green.duty_u16(65535)
        self.blue.duty_u16(65535)

# -------------------- Motion Function --------------------

def execute_movement(steps, direction=True):
    """Move all enabled harnesses with acceleration ramp."""
    dir_pin.value(direction)
    ramp = 20
    min_d = 1
    max_d = 10
    hold = steps - 2*ramp

    for i in range(ramp):
        step_pin.value(1)
        time.sleep_ms(max_d - int((max_d-min_d)*(i/ramp)))
        step_pin.value(0)
        time.sleep_ms(max_d - int((max_d-min_d)*(i/ramp)))

    for _ in range(hold):
        step_pin.value(1)
        time.sleep_ms(min_d)
        step_pin.value(0)
        time.sleep_ms(min_d)

    for i in range(ramp, 0, -1):
        step_pin.value(1)
        time.sleep_ms(max_d - int((max_d-min_d)*(i/ramp)))
        step_pin.value(0)
        time.sleep_ms(max_d - int((max_d-min_d)*(i/ramp)))

# -------------------- Initialization --------------------

# Create harness objects
harnesses = [Harness(i) for i in range(NUM_HARNESSES)]

# Clear shift register outputs
for h in harnesses:
    h.enable()
    h.disable()

# Initialize RGB LED
rgb_led = RGBLed(RGB_RED_PIN, RGB_GREEN_PIN, RGB_BLUE_PIN)

# Load WIF file
treadling = Treadling()
treadling.import_wif("sample.txt")

current_row = 0
rgb_led.set_color(*treadling.get_row_color(current_row))
disable_all_harnesses()

# Button edge states
execute_button_last = 1
prev_row_last = 1
next_row_last = 1

print("Starting on row", current_row+1, ", Color", treadling.get_row_color_index(current_row)+1)

lcd.clear()
lcd.move_to(0,0)
lcd.putstr("Row: {}".format(current_row + 1))
lcd.move_to(0,1)
lcd.putstr("Color: {}".format(treadling.get_row_color_index(current_row)+1))

# -------------------- Main Loop --------------------

while True:
    # Execute button
    execute_now = execute_pin.value()
    if execute_now == 0 and execute_button_last == 1:
        active_harnesses = treadling.get_row(current_row)
        enable_row_harnesses(active_harnesses)
        print("Raising row", current_row + 1)
        execute_movement(MOVEMENT_AMOUNT, True)
    elif execute_now == 1 and execute_button_last == 0:
        active_harnesses = treadling.get_row(current_row)
        enable_row_harnesses(active_harnesses)
        print("Lowering row", current_row + 1)
        execute_movement(MOVEMENT_AMOUNT, False)
        disable_all_harnesses()
    execute_button_last = execute_now

    # Next row button
    next_now = next_row_pin.value()
    if next_now == 0 and next_row_last == 1:
        current_row = (current_row + 1) % treadling.num_rows
        rgb_led.set_color(*treadling.get_row_color(current_row))
        print("Advanced to row", current_row + 1, ", Color", treadling.get_row_color_index(current_row)+1)
        time.sleep_ms(200)
    next_row_last = next_now

    # Previous row button
    prev_now = prev_row_pin.value()
    if prev_now == 0 and prev_row_last == 1:
        if current_row > 0:
            current_row -= 1
            rgb_led.set_color(*treadling.get_row_color(current_row))
            print("Back to row", current_row + 1, ", Color", treadling.get_row_color_index(current_row)+1)
        else:
            print("Already at first row")
        time.sleep_ms(200)
    prev_row_last = prev_now

    time.sleep_ms(10)