#!/usr/bin/env python3
"""Replace the HDL Coder `persistent isempty-init` table idiom with dist ROM IP.

Idiom (MATLAB `persistent T; if isempty(T), T = <const>; end`):

    wire  [..] P     [0:D-1];   +  D const assigns          (the constant image)
    wire  [..] T_1   [0:D-1];   generate: T_1[i]   = P[i];
    reg   [..] T_reg_reg [0:D-1];  always@(posedge clk/reset):
                                     reset -> 0; enb -> T_reg_reg <= T_reg_reg_next
    wire  [..] T_reg_reg_next [0:D-1];  generate: T[i] = T_reg_reg[i];
                                                  T_reg_reg_next[i] = TMP[i];
    reg   [..] TMP [0:D-1];     always@*: TMP = (S) ? T_1 : T;   // S = first-call flag

Net semantics: T reads 0 until the first (enb && S) tick, then the constant
image forever.  TMP reads (S || loaded) ? const : 0.

Replacement (bit- and cycle-true):
    dist_mem_gen ROM (content = P) per read site
    reg T__ipinited: reset->0, (enb && S)->1
    read T[e]   -> (T__ipinited ? ROM[e] : 0)
    read TMP[e] -> ((S || T__ipinited) ? ROM[e] : 0)
    all five arrays + three blocks deleted.

Usage: python initrom_to_ip.py <hdlsrc>   (operates on <hdlsrc>/rtl_ip copies;
run AFTER rtl_ip_scan + source copy [rom_to_ip.py --copy-only or manual],
BEFORE rom_to_ip's rewriting of any remaining plain ROMs, BEFORE namespace_ip.)
"""
import argparse, hashlib, json, os, re, shutil, sys

CONST_ASSIGN = re.compile(r"^\s*assign\s+(\w+)\[(\d+)\]\s*=\s*(\d+)'s?([bdh])([0-9a-fA-F_]+)\s*;\s*$", re.M)

def to_bits(radix, lit, width):
    lit = lit.replace('_', '')
    v = int(lit, {'b': 2, 'd': 10, 'h': 16}[radix]) & ((1 << width) - 1)
    return format(v, '0%db' % width)

def find_block(lines, i0, opener):
    """Return (start,end) line span of the statement block whose header line index
    is i0 (an `always`/`generate` line); tracks begin/end + generate/endgenerate."""
    depth, i = 0, i0
    started = False
    while i < len(lines):
        t = lines[i]
        depth += len(re.findall(r'\bbegin\b|\bgenerate\b|\bcase\b', t))
        depth -= len(re.findall(r'\bend\b(?!module|generate|case)|\bendgenerate\b|\bendcase\b', t))
        if depth > 0:
            started = True
        if started and depth <= 0:
            return i0, i
        i += 1
    sys.exit(f'unbalanced block at line {i0}: {lines[i0]}')

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('hdlsrc')
    ap.add_argument('outdir', help='FIXED output dir (e.g. <repo>/vivado/rtl_ip/<model>)')
    a = ap.parse_args()
    from ipflow_common import copy_sources
    hdlsrc = os.path.abspath(a.hdlsrc)
    rtl_ip = os.path.abspath(a.outdir)
    copy_sources(hdlsrc, rtl_ip)

    ipman_p = os.path.join(rtl_ip, 'ip_manifest.json')
    ipman = json.load(open(ipman_p)) if os.path.exists(ipman_p) else {'ips': []}
    coe_of_hash = {ip.get('content_hash'): ip['name'] for ip in ipman['ips'] if 'content_hash' in ip}

    for fn in [f for f in os.listdir(rtl_ip) if f.endswith('.v')]:
        path = os.path.join(rtl_ip, fn)
        txt = open(path, encoding='utf-8', errors='replace').read()

        # locate idiom instances: assign T_1[gv] = P[gv];  inside a generate loop
        pairs = []
        for m in re.finditer(r'assign\s+(\w+)_1\[(\w+)\]\s*=\s*(\w+)\[\2\]\s*;', txt):
            T, P = m.group(1), m.group(3)
            if f'{T}_reg_reg' in txt:
                pairs.append((T, P))
        if not pairs:
            continue
        print(f'[{fn}] isempty-init tables: {pairs}')

        for T, P in pairs:
            # ---- table content ----
            vals, W = {}, None
            for m in CONST_ASSIGN.finditer(txt):
                if m.group(1) == P:
                    W = W or len(to_bits(m.group(4), m.group(5), int(m.group(3))))
                    vals[int(m.group(2))] = to_bits(m.group(4), m.group(5), int(m.group(3)))
            D = len(vals)
            if D == 0 or set(vals) != set(range(D)):
                print(f'  SKIP {T}: const image {P} incomplete'); continue
            AW = max(1, (D - 1).bit_length())
            content = '\n'.join(vals[i] for i in range(D))
            h = hashlib.sha1((f'{W}x{D}\n' + content).encode()).hexdigest()[:10]
            if h not in coe_of_hash:
                ipname = f'rom_{W}x{D}_{h}'
                with open(os.path.join(rtl_ip, ipname + '.coe'), 'w') as f:
                    f.write('memory_initialization_radix=2;\nmemory_initialization_vector=\n')
                    f.write(',\n'.join(vals[i] for i in range(D)) + ';\n')
                ipman['ips'].append({'type': 'dist_mem_gen', 'name': ipname, 'width': W, 'depth': D,
                                     'coe': os.path.abspath(os.path.join(rtl_ip, ipname + '.coe')),
                                     'content_hash': h})
                coe_of_hash[h] = ipname
            ipname = coe_of_hash[h]

            lines = txt.split('\n')

            # ---- discover TMP name + guard signal from the mux always block ----
            mm = re.search(r'(\w+)\[\w+\]\s*=\s*%s_1\[\w+\]\s*;' % re.escape(T), txt)
            if not mm:
                print(f'  SKIP {T}: mux TMP not found'); continue
            TMP = mm.group(1)
            gm = None
            for i, ln in enumerate(lines):
                if re.search(r'\bif\s*\(\s*(\w+)\s*==\s*1\'b0\s*\)', ln) is not None:
                    # candidate: block must contain TMP writes
                    s2, e2 = find_block(lines, max(0, i - 1), 'always')
                    body = '\n'.join(lines[s2:e2 + 1])
                    if f'{TMP}[' in body and f'{T}_1[' in body:
                        gm = re.search(r'\bif\s*\(\s*(\w+)\s*==\s*1\'b0\s*\)', ln).group(1)
                        break
            if not gm:
                print(f'  SKIP {T}: guard signal not found'); continue
            S = gm

            # ---- delete the three blocks (reg process, copy generate, next generate, mux always) ----
            kill = set()
            for i, ln in enumerate(lines):
                if (re.search(r'begin\s*:\s*%s_reg_process' % re.escape(T), ln)
                        or re.search(r'begin\s*:\s*%s_1_gen' % re.escape(T), ln)
                        or re.search(r'begin\s*:\s*%s_reg_reg_next_gen' % re.escape(T), ln)):
                    j = i
                    while j > 0 and not re.search(r'^\s*(always|generate|genvar|\s*for)', lines[j]):
                        j -= 1
                    # generate blocks start at 'generate'/'genvar'; always at 'always'
                    while j > 0 and re.match(r'^\s*for\b', lines[j]):
                        j -= 1
                    if re.match(r'^\s*genvar', lines[j - 1] if j else ''):
                        j -= 1
                    s2, e2 = find_block(lines, j, None)
                    # extend to endgenerate if present
                    while e2 + 1 < len(lines) and re.match(r'^\s*endgenerate', lines[e2 + 1]):
                        e2 += 1
                    kill.update(range(s2, e2 + 1))
                if re.match(r'^\s*always\s*@\*', ln):
                    s2, e2 = find_block(lines, i, None)
                    body = '\n'.join(lines[s2:e2 + 1])
                    if f'{TMP}[' in body and f'{T}_1[' in body and S in body:
                        kill.update(range(s2, e2 + 1))
            # delete array declarations
            for i, ln in enumerate(lines):
                if re.search(r'\b(%s|%s_1|%s_reg_reg|%s_reg_reg_next|%s)\s*\[0:%d\]\s*;'
                             % tuple(map(re.escape, (T, T, T, T, TMP))) % () if False else
                             r'\b(' + '|'.join(map(re.escape, (T, T + '_1', T + '_reg_reg',
                                                               T + '_reg_reg_next', TMP))) + r')\s+\[0:' + str(D - 1) + r'\]\s*;', ln):
                    kill.add(i)
                if re.match(r'^\s*assign\s+%s\[\d+\]' % re.escape(P), ln):
                    kill.add(i)
                if re.search(r'\b%s\s+\[0:%d\]\s*;' % (re.escape(P), D - 1), ln):
                    kill.add(i)
            survivors = [ln for i, ln in enumerate(lines) if i not in kill]
            txt = '\n'.join(survivors)

            # ---- rewrite remaining real read sites ----
            inited = f'{T}__ipinited'
            inst, k = [], 0
            zeros = f"{W}'sb" + '0' * W
            for arr, gate in ((T, inited), (TMP, f'({S} || {inited})')):
                for m in list(re.finditer(re.escape(arr) + r'\[([^\]\n]+)\]', txt)):
                    expr = m.group(1)
                    if expr.isdigit():
                        continue
                    aq, qq = f'{T}_ipa{k}', f'{T}_ipq{k}'
                    inst.append(f'  wire [{AW-1}:0] {aq};\n  wire [{W-1}:0] {qq};\n'
                                f'  assign {aq} = {expr};\n'
                                f'  {ipname} u_{T}_ip{k} (.a({aq}), .spo({qq}));\n')
                    txt = txt.replace(f'{arr}[{expr}]', f'({gate} ? $signed({qq}) : {zeros})')
                    k += 1
            flag = (f'  reg {inited};\n'
                    f'  always @(posedge clk or posedge reset)\n'
                    f'    if (reset) {inited} <= 1\'b0;\n'
                    f'    else if (enb && ({S})) {inited} <= 1\'b1;\n')
            txt = txt.replace('endmodule',
                              f'  // ---- rtl-ip-replace(initrom): {T} -> {ipname} ({k} read site(s)) ----\n'
                              + flag + ''.join(inst) + '\nendmodule', 1)
            print(f'  {T} ({W}x{D}) -> {ipname}, {k} read site(s), guard={S}')

        open(path, 'w', encoding='utf-8').write(txt)

    json.dump(ipman, open(ipman_p, 'w'), indent=1)
    with open(os.path.join(rtl_ip, 'ip_manifest.tcl'), 'w') as f:
        for ip in ipman['ips']:
            f.write('lappend ips {%s}\n' % ' '.join(f'{k} {{{v}}}' for k, v in ip.items()
                                                    if k != 'content_hash'))
    print('ip_manifest updated')

if __name__ == '__main__':
    main()