from core import *

class AND(Chip):
    @staticmethod
    def logic(chip):
        if chip.pinIN() != (pow(2, chip.width) - 1) : chip.pinOUT(0)
        else : chip.pinOUT(1)
    def __init__(self, width):
        self.width = width
        self.datamask = slice(0, width)
        super().__init__(BitSet(width), BitSet(1))

class NAND(Chip):
    @staticmethod
    def logic(chip):
        if chip.pinIN() != (pow(2, chip.width) - 1) : chip.pinOUT(1)
        else : chip.pinOUT(0)
    def __init__(self, width):
        self.width = width
        self.datamask = slice(0, width)
        super().__init__(BitSet(width), BitSet(1))

class OR(Chip):
    @staticmethod
    def logic(chip):
        if chip.pinIN() != 0 : chip.pinOUT(1)
        else : chip.pinOUT(0)
    def __init__(self, width):
        self.width = width
        self.datamask = slice(0, width)
        super().__init__(BitSet(width), BitSet(1))

class NOR(Chip):
    @staticmethod
    def logic(chip):
        if chip.pinIN() == 0 : chip.pinOUT(1)
        else : chip.pinOUT(0)
    def __init__(self, width):
        self.width = width
        self.datamask = slice(0, width)
        super().__init__(BitSet(width), BitSet(1))

class NOT(Chip):
    @staticmethod
    def logic(chip):
        chip.pinOUT(int(str(chip.pinIN).replace("1", "l").replace("0", "1").replace("l", "0"), 2))
    def __init__(self, width):
        self.width = width
        self.datamask = slice(0, width)
        super().__init__(BitSet(width), BitSet(width))

class Timer(Chip):
    @staticmethod
    def logic(chip, step=False):
        if step : input("<step mode>:[on] press enter")
        else : print("<step mode>:[off]")
        if not chip[chip.hltmask] : 
            if chip[chip.tmask] : chip.pinOUT(1)
            else : chip.pinOUT(2)
        else : print("Timer HLT")
    def __init__(self):
        self.hltmask = 0
        self.tmask = 1
        super().__init__(BitSet(1), BitSet(2)(1), True)
    def update(self, *, step=False):
        self.logic(self, step)

class Register(Chip):
    @staticmethod
    def logic(chip):
        if chip.pinIN[chip.rstmask] : chip.data(0)
        if chip.pinIN[chip.clkmask] & chip.pinIN[chip.eimask] : chip.data(chip.pinIN[chip.datamask])
        if chip.pinIN[chip.eomask] : chip.pinOUT(chip.data())
        else : chip.pinOUT(0)
    def __init__(self, width):
        self.width = width
        self.rstmask = 0
        self.clkmask = 1
        self.eimask  = 2
        self.eomask  = 3
        self.datamask = slice(4, 4+width)
        self.data = BitSet(width)
        super().__init__(BitSet(width+4), BitSet(width), True)
    def input(self, rst, clk, ei, eo, data):
        self.pinIN[self.rstmask]  = rst
        self.pinIN[self.clkmask]  = clk
        self.pinIN[self.eimask]   = ei
        self.pinIN[self.eomask]   = eo
        self.pinIN[self.datamask] = data
        self.logic(self)

class Adder(Chip):
    @staticmethod
    def logic(chip):
        if chip.pinIN[chip.submask] : chip.data(chip.pinIN[chip.data1mask] + (pow(2, chip.width) - chip.pinIN[chip.data2mask]))
        else : chip.data(chip.pinIN[chip.data1mask] + chip.pinIN[chip.data2mask])
        if chip.pinIN[chip.eomask] : chip.pinOUT(chip.data())
        else : chip.pinOUT(0)
    def __init__(self, width):
        self.width = width
        self.eomask    = 0
        self.submask   = 1
        self.data1mask = slice(2, 2+width)
        self.data2mask = slice(2+width, 2+width*2)
        self.data = BitSet(width + 1)
        super().__init__(BitSet(width*2+2), BitSet(width+1))
    def input(self, eo, sub, data1, data2):
        self.pinIN[self.eomask]    = eo
        self.pinIN[self.submask]   = sub
        self.pinIN[self.data1mask] = data1
        self.pinIN[self.data2mask] = data2
        self.logic(self)
    def show(self):
        print(self.data() if self.data()<128 else self.data()-256)

class Memory(Chip):
    @staticmethod
    def logic(chip:BitSet):
        if chip.pinIN[chip.aimask] : chip.sel_addr = chip.pinIN[chip.addrmask]
        if chip.pinIN[chip.momask] : chip.pinOUT(chip.data_pool[chip.sel_addr]())
        else : chip.pinOUT(0)
        if chip.pinIN[chip.mimask] : chip.data_pool[chip.sel_addr](chip.pinIN[chip.datamask])
    def __init__(self, addr_width, data_width):
        import random
        self.a_width = addr_width
        self.d_width = data_width
        self.mimask = 0
        self.momask = 1
        self.aimask = 2
        self.addrmask = slice(3, 3+addr_width)
        self.datamask = slice(3+addr_width, 3+addr_width+data_width)
        self.sel_addr = 0
        self.data_pool = [BitSet(data_width)(random.randint(1, pow(2,data_width))-1) for i in range(pow(2, addr_width)) ]
        super().__init__(BitSet(addr_width+data_width+3), BitSet(data_width))
    def input(self, ei, eo, addr, data):
        self.pinIN[self.mimask]   = ei
        self.pinIN[self.momask]   = eo
        self.pinIN[self.addrmask] = addr
        self.pinIN[self.datamask] = data
        self.logic(self)
    def load_mem(self, addr, data):
        self.data_pool[addr](data)
    def show(self, *, key=None):
        print(self.__repr__())
        if key == None : key = slice(0, pow(2, self.a_width), 1)
        elif isinstance(key, int) : key = slice(key, key+1, 1)
        for data in self.data_pool[key]:
            print((" {:>3d} : {:0>%db} > %s"%(self.a_width, data)).format(self.data_pool.index(data), self.data_pool.index(data)))

class Counter(Chip):
    @staticmethod
    def logic(chip):
        if chip.pinIN[chip.rstmask] : 
            chip.pinOUT(0)
            chip.data(0)
        else :
            if chip.pinIN[chip.comask] : chip.pinOUT(chip.data())
            else : chip.pinOUT(0)
            if chip.pinIN[chip.cimask] : chip.data(chip.pinIN[chip.datamask])
            elif chip.pinIN[chip.enmask] :
                if chip.pinIN[chip.clkmask] & (not chip.lastclk):
                    chip.data((chip.data()+1)&(pow(2, chip.width)-1))
        chip.lastclk = chip.pinIN[chip.clkmask]
    def __init__(self, width):
        self.width = width
        self.rstmask = 0
        self.clkmask = 1
        self.enmask = 2
        self.cimask = 3
        self.comask = 4
        self.datamask = slice(5, 5+width)
        self.lastclk = 0
        self.data = BitSet(width)
        super().__init__(BitSet(5+width), BitSet(width), True)
    def input(self, rst, clk, en, ci, data):
        self.pinIN[self.rstmask]  = rst
        self.pinIN[self.clkmask]  = clk
        self.pinIN[self.enmask]   = en
        self.pinIN[self.cimask]   = ci
        self.pinIN[self.comask]   = co
        self.pinIN[self.datamask] = data
        self.logic(self)

class PROM(Chip):
    @staticmethod
    def logic(chip):
        chip.pinOUT(chip.rom[chip.pinIN()])
    def __init__(self, addr_width, data_width, rom):
        self.a_width = addr_width
        self.d_width = data_width
        self.addrmask = slice(0, addr_width)
        self.rom = rom
        super().__init__(BitSet(addr_width), BitSet(data_width)(rom[0]))
    def input(self, addr):
        self.pinIN[self.addrmask] = addr
        self.logic(self)

class Encoder(Chip):
    @staticmethod
    def logic(chip):
        for i in range(pow(2, chip.width)):
            if chip.pinIN[i] == 1 :
                chip.pinOUT(i)
                break
    def __init__(self, width):
        self.width = width
        self.datamask = slice(0, pow(2, width))
        super().__init__(BitSet(pow(2, width)), BitSet(width))
    def input(self, data):
        self.pinIN[self.datamask] = data
        self.logic(self)

class Decoder(Chip):
    @staticmethod
    def logic(chip):
        chip.pinOUT(pow(2, chip.pinIN()))
    def __init__(self, width):
        self.width = width
        self.datamask = slice(0, width)
        super().__init__(BitSet(width), BitSet(pow(2, width)))
    def input(self, data):
        self.pinIN[self.datamask] = data
        self.logic(self)