# ============================================================
# NEC IR Transmitter Demo for RP2040 (Raspberry Pi Pico)
# Integrated IR base + RP2_RMT + NEC protocol
# ============================================================

from sys import platform
from micropython import const
from array import array
from time import ticks_us, ticks_diff, sleep_ms
from machine import Pin, PWM
import rp2

# ------------------------------------------------------------
# 平台判断
# ------------------------------------------------------------
ESP32 = platform == 'esp32'
RP2 = platform == 'rp2'

assert RP2, "This demo is intended to run on RP2040 (rp2)"

# ------------------------------------------------------------
# RP2040 PIO RMT 实现（来自 rp2_rmt.py）
# ------------------------------------------------------------

@rp2.asm_pio(autopull=True, pull_thresh=32) # 每次从 FIFO 取 32 位，正好一个时间值
def irqtrain(): # 中断节拍器
    wrap_target() 
    out(x, 32)  # 从 FIFO 里取 32 位数据，放进 X 寄存器
    irq(rel(0))
    label("loop")
    jmp(x_dec, "loop")
    wrap()


class DummyPWM:
    def duty_u16(self, _):
        pass


class RP2_RMT:
    def __init__(self, pin_pulse=None, carrier=None, sm_no=0, sm_freq=1_000_000):
        if carrier is None:
            self.pwm = DummyPWM()
            self.duty = (0, 0)
        else:
            pin_car, freq, duty = carrier
            self.pwm = PWM(pin_car)
            self.pwm.freq(freq)
            self.pwm.duty_u16(0)
            self.duty = (int(0xFFFF * duty // 100), 0)

        self.sm = rp2.StateMachine(sm_no, irqtrain, freq=sm_freq)
        self.apt = 0
        self.arr = None
        self.ict = None
        self.icm = 0
        self.reps = 0

        rp2.PIO(0).irq(
            handler=self._cb,
            trigger=1 << (sm_no + 8),
            hard=True
        )

    def _cb(self, _):
        if self.ict is not None:
            self.pwm.duty_u16(self.duty[self.ict & 1])
            self.ict += 1

            if d := self.arr[self.apt]:
                self.sm.put(d)
                self.apt += 1
            else:
                if self.reps != 1:
                    if self.reps:
                        self.reps -= 1
                    self.sm.put(self.arr[0])
                    self.apt = 1
                    self.ict = 1

    def send(self, ar, reps=1):
        self.sm.active(0)
        self.reps = reps
        ar[-1] = 0

        for i, v in enumerate(ar):
            if v == 0:
                break

        self.icm = i
        mv = memoryview(ar)
        self.sm.put(mv[:min(i, 4)])

        self.arr = ar
        self.apt = min(i, 4)
        self.ict = 0
        self.sm.active(1)

    def busy(self):
        return self.ict is not None and self.ict < self.icm


# ------------------------------------------------------------
# IR 抽象基类
# ------------------------------------------------------------

STOP = const(0)

class IR:
    timeit = False

    def __init__(self, pin, cfreq, asize, duty, verbose):
        self._rmt = RP2_RMT(
            pin_pulse=None,
            carrier=(pin, cfreq, duty)
        )

        self._arr = array('H', 0 for _ in range(asize + 1))
        self._mva = memoryview(self._arr)

        self.verbose = verbose
        self.aptr = 0

    def busy(self):
        return self._rmt.busy()

    def append(self, *times):
        for t in times:
            self._arr[self.aptr] = t
            self.aptr += 1
            self.verbose and print("append:", t)

    def trigger(self):
        self.append(STOP)
        self._rmt.send(self._arr)

    def transmit(self, addr, data, toggle=0, validate=False):
        while self.busy():
            pass

        t0 = ticks_us()
        self.aptr = 0
        self.tx(addr, data, toggle)
        self.trigger()

        if self.timeit:
            print("TX time:", ticks_diff(ticks_us(), t0), "us")


# ------------------------------------------------------------
# NEC 协议实现
# ------------------------------------------------------------

_TBURST = const(563)
_T_ONE = const(1687)

class NEC(IR):
    valid = (0xffff, 0xff, 0)
    samsung = False

    def __init__(self, pin, freq=38000, verbose=False):
        super().__init__(pin, freq, 68, 33, verbose)

    def _bit(self, b):
        self.append(_TBURST, _T_ONE if b else _TBURST)

    def tx(self, addr, data, _):
        self.append(9000, 4500)

        if addr < 256:
            addr |= ((addr ^ 0xff) << 8)

        for _ in range(16):
            self._bit(addr & 1)
            addr >>= 1

        data |= ((data ^ 0xff) << 8)

        for _ in range(16):
            self._bit(data & 1)
            data >>= 1

        self.append(_TBURST)

    def repeat(self):
        self.aptr = 0
        self.append(9000, 2250, _TBURST)
        self.trigger()


# ------------------------------------------------------------
# Demo 主程序
# ------------------------------------------------------------

IR_PIN = 17      # IR LED
LED_PIN = 20     # Pico 板载 LED

ir_led = Pin(IR_PIN, Pin.OUT, value=0)
led = Pin(LED_PIN, Pin.OUT)

ir = NEC(ir_led)

ADDR = 0x10
DATA = 0x0B

print("RP2040 NEC IR transmitter running...")
print("ADDR =", hex(ADDR), "DATA =", hex(DATA))

while True:
    ir.transmit(ADDR, DATA)
    led.toggle()
    sleep_ms(500)