from __future__ import absolute_import
from __future__ import print_function
import sys
import os
import math
from functools import reduce

# the next line can be removed after installation
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from veriloggen import *
import veriloggen.dataflow as dataflow
import veriloggen.types.ram as ram
import veriloggen.types.fixed as fixed

def stencil(coe, data):
    data = map(lambda x,y: x*y, data, coe)
    rslt = reduce(lambda x,y: x+y, data)
    return rslt
    
def mkStencilPipeline2D(coe=None, size=3, width=32, point=16):
    # size-port stream inputs
    iports = [ dataflow.Variable('idata%d'%i, valid='ivalid%d'%i,
                                 width=width, point=point, signed=True)
               for i in range(size) ]

    if coe is None:
        coe = [ [ dataflow.Constant(1.0/(1.0*size*size), point=point) for i in range(size) ]
                for j in range(size) ]
        
    # source data
    data = [ [ d.prev(j) for j in range(size) ] for d in iports ]
    
    # from 2D list to 1D list
    data_list = []
    coe_list = []
    for d, c in zip(data, coe):
        data_list.extend(d)
        coe_list.extend(c)

    # computation by calling standard method
    rslt = stencil(coe_list, data_list)
    
    rslt.output('odata', valid='ovalid')
    
    df = dataflow.Dataflow(rslt)
    m = df.to_module('stencil_pipeline_2d')

    #try:
    #    df.draw_graph()
    #except:
    #    print('Dataflow graph could not be generated.', file=sys.stderr)

    return m

def mkStencil(n=16, size=3, datawidth=32, point=16, coe_test=False):
    m = Module('stencil')

    addrwidth = int(math.log(n, 2))

    clk = m.Input('CLK')
    rst = m.Input('RST')

    start = m.Input('start')
    busy = m.OutputReg('busy', initval=0)
    
    done = m.TmpReg(initval=0)

    # external RAM I/F
    ext_src_rams = [ ram.RAMSlaveInterface(m, 'ext_src_ram%d' % i,
                                              datawidth=datawidth, addrwidth=addrwidth)
                      for i in range(size) ]
    ext_dst_ram = ram.RAMSlaveInterface(m, 'ext_dst_ram',
                                           datawidth=datawidth, addrwidth=addrwidth)
    
    # RAM
    addrwidth = int(math.log(n, 2)) * 2
    
    src_rams = [ ram.SyncRAMManager(m, 'src_ram%d' % i, clk, rst,
                                    datawidth=datawidth, addrwidth=addrwidth, numports=2)
                  for i in range(size) ]
                            
    dst_ram = ram.SyncRAMManager(m, 'dst_ram', clk, rst,
                                 datawidth=datawidth, addrwidth=addrwidth, numports=2)

    # connect RAM I/Fs
    for src_ram, ext_src_ram in zip(src_rams, ext_src_rams):
        src_ram[1].connect(ext_src_ram)

    dst_ram[1].connect(ext_dst_ram)

    # read FSM
    read_fsm = FSM(m, 'read_fsm', clk, rst)
    read_count = m.Reg('read_count', 32, initval=0)
    read_addr = m.Reg('read_addr', 32, initval=0)


    read_fsm(
        read_addr(0),
        read_count(0),
        busy(0)
    )

    read_fsm.If(start)(
        busy(1)
    )

    read_fsm.Then().goto_next()

    read_fsm(
        read_addr.inc(),
        read_count.inc()
    )
        
    idata = []
    ivalid = []
    for i, src_ram in enumerate(src_rams):
        src_ram.disable_write(0)
        rdata, rvalid = src_ram.read(0, read_addr, read_fsm)
        idata.append(rdata)
        ivalid.append(rvalid)

    read_fsm.If(read_count == n - 1)(
        read_addr(0),
        read_count(0)
    )

    read_fsm.Then().goto_next()

    read_fsm.If(done)(
        busy(0)
    )

    read_fsm.Then().goto_init()

    read_fsm.make_always()


    # instance
    odata = m.Wire('odata', datawidth)
    ovalid = m.Wire('ovalid')
    
    ports = []
    ports.append( ('CLK', clk) )
    ports.append( ('RST', rst) )    
    
    for i, (d, v) in enumerate(zip(idata, ivalid)):
        ports.append( ('idata%d' % i, d) )
        ports.append( ('ivalid%d' % i, v) )

    ports.append( ('odata', odata) )
    ports.append( ('ovalid', ovalid) )

    coe = None
    if coe_test:
        coe = [ [ dataflow.Constant(1, point=point) for i in range(size) ]
                for j in range(size) ]
        point = 0
        
    st = mkStencilPipeline2D(size=3, width=datawidth, point=point, coe=coe)
    m.Instance(st, 'inst_stencil', ports=ports)

    skip_offset = int(math.floor(size/2))
    
    # write FSM
    write_fsm = FSM(m, 'write_fsm', clk, rst)
    write_count = m.Reg('write_count', 32, initval=0)
    write_addr = m.Reg('write_addr', 32, initval=skip_offset)

    write_fsm(
        done(0)
    )
    
    write_fsm.If(Ands(ovalid, write_count > skip_offset))(
        write_addr.inc()
    )
    
    dst_ram.write(0, write_addr, odata, write_fsm.then)
    
    write_fsm.If(ovalid)(
        write_count.inc(),
    )
    
    write_fsm.If(write_count == n)(
        write_count(0),
        write_addr(skip_offset),
        done(1)
    )
    write_fsm.Then().goto_init()

    write_fsm.make_always()

    return m

def mkTest(n=16, size=3, datawidth=32, point=16, coe_test=False):
    if coe_test:
        point = 0
    
    m = Module('test')

    addrwidth = int(math.log(n, 2))

    main = mkStencil(n, size, datawidth, point, coe_test)

    params = m.copy_params(main)
    ports = m.copy_sim_ports(main)
    
    clk = ports['CLK']
    rst = ports['RST']

    start = ports['start']
    busy = ports['busy']
    
    uut = m.Instance(main, 'uut',
                     params=m.connect_params(main),
                     ports=m.connect_ports(main))

    reset_done = m.Reg('reset_done', initval=0)
    reset_stmt = []
    reset_stmt.append( reset_done(0) )
    reset_stmt.append( start(0) )

    
    # src RAM
    for i in range(3):
        addr = ports['ext_src_ram%d_addr' % i]
        rdata = ports['ext_src_ram%d_rdata' % i]
        wdata = ports['ext_src_ram%d_wdata' % i]
        wenable = ports['ext_src_ram%d_wenable' % i]
        reset_stmt.append(addr(0))
        reset_stmt.append(wdata(0))
        reset_stmt.append(wenable(0))

    # dst RAM
    addr = ports['ext_dst_ram_addr']
    rdata = ports['ext_dst_ram_rdata']
    wdata = ports['ext_dst_ram_wdata']
    wenable = ports['ext_dst_ram_wenable']
    reset_stmt.append(addr(2))
    reset_stmt.append(wdata(0))
    reset_stmt.append(wenable(0))
    

    simulation.setup_waveform(m, uut)
    simulation.setup_clock(m, clk, hperiod=5)
    init = simulation.setup_reset(m, rst, reset_stmt, period=100)

    nclk = simulation.next_clock

    init.add(
        Delay(1000),
        reset_done(1),
        nclk(clk),
        Delay(100000),
        Systask('finish'),
    )

    fsm = FSM(m, 'fsm', clk, rst)
    
    fsm.goto_next(cond=reset_done)

    for i in range(3):
        addr = ports['ext_src_ram%d_addr' % i]
        fsm.add( addr(-1) )

    fsm.goto_next()

    for i in range(3):
        addr = ports['ext_src_ram%d_addr' % i]
        rdata = ports['ext_src_ram%d_rdata' % i]
        wdata = ports['ext_src_ram%d_wdata' % i]
        wenable = ports['ext_src_ram%d_wenable' % i]
        next_addr = (addr+1) % (n*n)
        fsm.add( addr.inc() )
        fsm.add( wdata(fixed.FixedConst(90, point).raw) )
        fsm.add( wenable(1) )
        fsm.add( wenable(0), cond=AndList(wenable, addr==2**addrwidth-1) )

    fsm.goto_next(cond=AndList(wenable, ports['ext_src_ram0_addr']==2**addrwidth-1))
    
    fsm.goto_next(cond=Not(busy))
    
    fsm.add( start(1) )
    fsm.add( start(0), delay=1 )
    fsm.goto_next()
    
    fsm.goto_next(cond=busy)

    fsm.goto_next(cond=Not(busy))

    fsm.add( Systask('finish') )

    fsm.make_always()

    return m

if __name__ == '__main__':
    n = 16
    #test = mkTest(n, coe_test=True)
    test = mkTest(n)
    verilog = test.to_verilog('tmp.v')
    #print(verilog)

    # run simulator (Icarus Verilog)
    sim = simulation.Simulator(test)
    rslt = sim.run()
    print(rslt)
