from math import log2, ceil
import os
import yaml
import sys
import shutil

parent_dir = os.path.dirname(os.path.abspath(__file__))
gen_dir = os.path.join(parent_dir, "../gen")
sw_dir = os.path.join(gen_dir,"sw")
init_dir = os.path.join(sw_dir,"init")
c_dir = os.path.join(sw_dir,"include")
v_dir = os.path.join(gen_dir,"rtl")

class CorevusConfig():

    def __init__(self,v_dwidth,i_cache_kbytes,d_cache_kbytes,crf_bytes):

        self.v_dwidth = v_dwidth
        self.vword_bytes = int(self.v_dwidth/8)
        self.vword_words = int(self.v_dwidth/32)
        self.vword_wordsel_bits = ceil(log2(self.vword_words))

        self.i_cache_bytes = i_cache_kbytes*1024
        self.d_cache_bytes = d_cache_kbytes*1024
        self.c_bytes = crf_bytes

        self.i_depth = self.i_cache_bytes/4
        self.i_awidth = log2(self.i_depth)

        self.d_depth = self.d_cache_bytes/self.vword_bytes
        self.d_awidth = log2(self.d_depth)
        self.d_bytesel = log2(self.vword_bytes)

        self.c_depth = self.c_bytes/self.vword_bytes
        self.c_awidth = log2(self.c_depth)

        self.ms_msb = 0
        self.ms_lsb = 0
        self.ms_msb_vword = 0
        self.ms_lsb_vword = 0

        self.i_cache_base_addr = 0
        self.d_cache_base_addr = 0
        self.c_reg_base_addr = 0

def usage():

    print ("Usage: python parse_config.py <path-to-config.yml>");
    exit()

def partition_memory(corevus_config):

    c_size_bytes = corevus_config.c_bytes
    i_size_bytes = corevus_config.i_cache_bytes
    d_size_bytes = corevus_config.d_cache_bytes

    req_c_addr_bits_b = log2(c_size_bytes) 
    req_i_addr_bits_b = log2(i_size_bytes)
    req_d_addr_bits_b = log2(d_size_bytes)

    req_max_addr_bits_b = max(req_d_addr_bits_b,max(req_c_addr_bits_b,req_i_addr_bits_b))

    corevus_config.ms_msb = req_max_addr_bits_b+1-2
    corevus_config.ms_lsb = req_max_addr_bits_b-2

    corevus_config.ms_msb_vword = corevus_config.ms_msb-corevus_config.vword_wordsel_bits
    corevus_config.ms_lsb_vword = corevus_config.ms_lsb-corevus_config.vword_wordsel_bits

    corevus_config.c_reg_base_addr = 0
    corevus_config.i_cache_base_addr = 2**(corevus_config.ms_lsb+2)
    corevus_config.d_cache_base_addr = 2**(corevus_config.ms_lsb+2)*2
    corevus_config.ref_phony_addr = 2**(corevus_config.ms_lsb+2)*3

def parse_yaml(yaml_raw):

    v_width = yaml_raw['vector_config']['dwidth']

    i_cache_kbytes = yaml_raw['local_memory']['icache_kbytes']
    d_cache_kbytes = yaml_raw['local_memory']['dcache_kbytes']
    crf_bytes = yaml_raw['local_memory']['crf_bytes']
    
    return CorevusConfig(v_width,i_cache_kbytes,d_cache_kbytes,crf_bytes)
    
def clean_gen():

    try:
        shutil.rmtree(gen_dir)
    except:
        pass
    os.mkdir(gen_dir)
    os.mkdir(v_dir)
    os.mkdir(sw_dir)
    os.mkdir(c_dir)
    os.mkdir(init_dir)

def gen_header(corevus_config):

    text = "#ifndef COREVUS_GEN_H\n#define COREVUS_GEN_H\n"
    text += "#define COREVUS_LOCAL_IMEM_BASE_ADDR "
    text += str(int(corevus_config.i_cache_base_addr))
    text += "\n"
    text += "#define COREVUS_LOCAL_DMEM_BASE_ADDR "
    text += str(int(corevus_config.d_cache_base_addr))
    text += "\n"
    text += "#define COREVUS_LOCAL_CRF_BASE_ADDR "
    text += str(int(corevus_config.c_reg_base_addr))
    text += "\n"

    text += "#endif\n"

    f = open(c_dir+"/corevus_gen.h","w")
    f.write(text)
    f.close()

def gen_params(corevus_config):

    text = "localparam MS_MSB = "
    text += str(int(corevus_config.ms_msb)) + ";\n"
    text += "localparam MS_LSB = "
    text += str(int(corevus_config.ms_lsb)) + ";\n"
    text += "localparam MS_MSB_DCACHE = "
    text += str(int(corevus_config.ms_msb_vword)) + ";\n"
    text += "localparam MS_LSB_DCACHE = "
    text += str(int(corevus_config.ms_lsb_vword)) + ";\n"
    text += "localparam DCACHE_BE_LEN = " 
    text += str(int(corevus_config.vword_bytes)) + ";\n"
    text += "localparam DCACHE_AWIDTH = "
    text += str(int(corevus_config.d_awidth)) + ";\n"
    text += "localparam DCACHE_DEPTH = "
    text += str(int(corevus_config.d_depth)) + ";\n"
    text += "localparam DCACHE_WE_LEN = "
    text += str(int(corevus_config.vword_wordsel_bits)) + ";\n"
    text += "localparam ICACHE_AWIDTH = "
    text += str(int(corevus_config.i_awidth)) + ";\n"
    text += "localparam ICACHE_DEPTH = "
    text += str(int(corevus_config.i_depth)) + ";\n"
    text += "localparam DCACHE_DWIDTH = "
    text += str(int(corevus_config.v_dwidth)) + ";\n"
    text += "localparam DCACHE_VWORD_AWIDTH=30-DCACHE_WE_LEN;\n"


    text += "localparam CRF_AWIDTH = "
    text += str(int(corevus_config.c_awidth)) + ";\n"
    text += "localparam CRF_DEPTH = "
    text += str(int(corevus_config.c_depth)) + ";\n"

    f = open(v_dir+"/params.vinc","w")
    f.write(text)
    f.close()

    text = "localparam IMEM_BASE_ADDR = 32'd" + str(int(corevus_config.i_cache_base_addr)) + ";\n"
    text += "localparam DMEM_BASE_ADDR = 32'd" + str(int(corevus_config.d_cache_base_addr)) + ";\n"
    text += "localparam CRF_BASE_ADDR = 32'd" + str(int(corevus_config.c_reg_base_addr)) + ";\n"
    text += "localparam IMEM_SIZE_BYTES = 32'd" + str(int(corevus_config.i_cache_bytes)) + ";\n"
    text += "localparam DMEM_SIZE_BYTES = 32'd" + str(int(corevus_config.d_cache_bytes)) + ";\n"
    text += "localparam CRF_SIZE_BYTES = 32'd" + str(int(corevus_config.c_bytes)) + ";\n"
    
    f = open(v_dir+"/reg_addr.vinc","w")
    f.write(text)
    f.close()

def gen_extract_mux(sel, sel_width, n_entries, stride, lhs, rhs):

    text = "    always_comb begin\n"
    text += "        case (" + str(sel) + ")\n"

    for i in range(0,n_entries):
        text += "             " + str(sel_width) + "'d"
        text += str(i) + ": " + lhs + " = " + rhs + "[" + str(stride) +"*" + str(i+1) + "-1:" + str(stride) + "*" + str(i) + "];\n"
    text += "            default: " + lhs + " = " + str(stride) + "'d0;\n"
    text += "        endcase\n"
    text += "    end\n\n"

    return text

def gen_pack_mux(sel, sel_width, n_entries, stride, lhs, rhs):

    text = "    always_comb begin\n"
    text += "        case (" + str(sel) + ")\n"

    for i in range(0, n_entries):
        pad_lh = (stride*n_entries)-stride*(i+1)
        pad_rh = stride*i
        text += "            " + str(sel_width) + "'d"
        text += str(i) + ": " + lhs + " = {"

        if (pad_lh != 0):
            text += str(pad_lh) + "'d0,"

        text += rhs

        if (pad_rh != 0):
            text += "," + str(pad_rh) + "'d0"

        text += "};\n"

    text += "            default: " + lhs + " = " + str(stride*n_entries) + "'d0;\n"
    text += "        endcase\n"
    text += "    end\n\n"

    return text

def gen_fab_dcache_mux(corevus_config):

    text = gen_extract_mux("r_dcache_wsel_fab_del", corevus_config.vword_wordsel_bits, corevus_config.vword_words, 32, "w_dcache_rdata_fab_out", "w_dcache_rdata_fab")
    text += gen_pack_mux("w_dcache_wsel_fab", corevus_config.vword_wordsel_bits, corevus_config.vword_words, 32, "w_dcache_wdata_fab", "i_fabric_wdata")
    text += gen_pack_mux("w_dcache_wsel_fab", corevus_config.vword_wordsel_bits, corevus_config.vword_words, 4, "w_dcache_wbe_fab", "{4{i_fabric_we}}")

    f = open(v_dir+"/fab_dcache_mux.vinc","w")
    f.write(text)
    f.close()

def gen_fab_crf_mux(corevus_config):
    text = gen_extract_mux("r_crf_wsel_fab", corevus_config.vword_wordsel_bits, corevus_config.vword_words, 32, "o_fabric_rdata", "r_crf_rdata_fab")
    text += gen_pack_mux("w_crf_wsel_fab", corevus_config.vword_wordsel_bits, corevus_config.vword_words, 32, "w_crf_wdata_fab", "i_fabric_wdata")
    text += gen_pack_mux("w_crf_wsel_fab", corevus_config.vword_wordsel_bits, corevus_config.vword_words, 4, "w_crf_wbe_fab", "{4{i_fabric_we}}")

    f = open(v_dir+"/fab_crf_mux.vinc","w")
    f.write(text)
    f.close()


def gen_lsu_mux(corevus_config):

    text = gen_extract_mux("r_word_sel_WB", corevus_config.vword_wordsel_bits, corevus_config.vword_words, 32, "w_dmem_rdata", "w_dmem_vword_rdata")
    text += gen_pack_mux("w_word_sel_EX1", corevus_config.vword_wordsel_bits, corevus_config.vword_words, 32, "w_dmem_vword_wdata", "w_dmem_wdata")
    text += gen_pack_mux("w_word_sel_EX1", corevus_config.vword_wordsel_bits, corevus_config.vword_words, 4, "w_dmem_vword_wbe", "w_dmem_wbe")

    f = open(v_dir+"/lsu_mux.vinc","w")
    f.write(text)
    f.close()


def gen_crf_reset(corevus_config):

    text = ""

    for i in range(0,int(corevus_config.c_depth)):
        text += "            regs["+str(i)+"] <= "+str(corevus_config.v_dwidth)+"'d0;\n"

    f = open(v_dir+"/crf_reset.vinc","w")
    f.write(text)
    f.close()

def gen_crf_core_wbe_assign(corevus_config):
    text = ""
    for i in range(0,corevus_config.vword_bytes):
        bit_sel = "[8*" + str(i+1) + "-1:8*" + str(i) + "]"
        text += "if (i_core_we[" + str(i) + "]) begin\n"
        text += "    regs[i_core_adx]" + bit_sel
        text += " <= i_core_wdata" + bit_sel + ";\n"
        text += "end\n"

    f = open(v_dir+"/crf_core_wbe_assign.vinc","w")
    f.write(text)
    f.close()

def gen_crf_fab_wbe_assign(corevus_config):
    text = ""
    for i in range(0,corevus_config.vword_bytes):
        bit_sel = "[8*" + str(i+1) + "-1:8*" + str(i) + "]"
        text += "if (w_crf_wbe_fab[" + str(i) + "]) begin\n"
        text += "    regs[w_crf_adx_fab]" + bit_sel
        text += " <= w_crf_wdata_fab" + bit_sel + ";\n"
        text += "end\n"

    f = open(v_dir+"/crf_fab_wbe_assign.vinc","w")
    f.write(text)
    f.close()

def gen_linker_script(corevus_config):

    text_base_hex = hex(int(corevus_config.i_cache_base_addr))
    data_base_hex = hex(int(corevus_config.d_cache_base_addr))
    ref_base_hex = hex(int(corevus_config.ref_phony_addr));

    text = "OUTPUT_FORMAT(\"elf32-littleriscv\", \"elf32-littleriscv\",\n"
    text += "\"elf32-littleriscv\")\n"
    text += "OUTPUT_ARCH(riscv)\n"
    text += "ENTRY(_start)\n"
    text += "SECTIONS\n"
    text += "{\n"
    text += ". = " + text_base_hex + ";\n"
    text += ".text : { *(.text) }\n"
    text += ". = " + data_base_hex + ";\n"
    text += ".data : { *(.data) }\n"
    text += ".bss : { *(.bss) }\n"
    text += ". = " + ref_base_hex + ";\n"
    text += ".ref : { *(.ref) }\n"
    text += "}\n"

    f = open(init_dir+"/linker_script","w")
    f.write(text)
    f.close()

def gen_init_asm(corevus_config):

    text = ".globl _start\n"
    text += "_start:\n"
    text += "li sp," + str(int(corevus_config.d_cache_base_addr+corevus_config.d_cache_bytes)) + "\n"
    text += "jal ra,main\n"

    f = open(init_dir+"/init.s","w")
    f.write(text)
    f.close()


if __name__ == "__main__":

    if len(sys.argv)!=2:
        usage()

    if not os.path.isfile(sys.argv[1]):
        usage()


    file_stream = open(sys.argv[1])
    spec_raw = yaml.load(file_stream,Loader=yaml.FullLoader)
    corevus_config = parse_yaml(spec_raw)
    partition_memory(corevus_config)
    clean_gen()
    gen_linker_script(corevus_config)
    gen_header(corevus_config)
    gen_params(corevus_config)
    gen_crf_reset(corevus_config)
    gen_crf_core_wbe_assign(corevus_config)
    gen_crf_fab_wbe_assign(corevus_config)
    gen_fab_dcache_mux(corevus_config)
    gen_fab_crf_mux(corevus_config)
    gen_lsu_mux(corevus_config)
    gen_init_asm(corevus_config)

