#!/usr/bin/env python3
"""Internal RTL Diagnostic-30 Verilog benchmark harness.

Inspired by VerilogEval/RTLLM-style evaluation:
- specification-to-RTL prompts
- exact module/port requirements
- compile with iverilog -g2012
- functional simulation with vvp when testbench exists
- category breakdown
- pass@k support via stochastic sampling
- saves generated RTL and JSONL records

Usage examples:
  python src/paper_style_verilog_benchmark.py --list-suites
  python src/paper_style_verilog_benchmark.py --suite full --model-name base --out internal_rtl_diagnostic_results/base
  python src/paper_style_verilog_benchmark.py --suite full --model-name v3 --adapter adapter_v3_functional --out internal_rtl_diagnostic_results/v3
"""
import argparse, json, math, os, random, re, subprocess, tempfile, time
from pathlib import Path

# Heavy ML imports are loaded lazily in load_model()/gen() so --list-suites works
# on machines without torch/transformers installed.

BASE = "Qwen/Qwen2.5-Coder-7B-Instruct"
SYSTEM = "Return only complete synthesizable Verilog code. No explanation. Use exact requested module name and ports."

TASKS = []
def task(tid, cat, prompt, tb=None, compile_only=False):
    TASKS.append({"id":tid,"category":cat,"prompt":prompt,"testbench":tb,"compile_only":compile_only})

# --- Basic combinational ---
task('comb_half_adder','basic_comb','Implement module half_adder(input a, input b, output sum, output carry).',
"module tb;reg a,b;wire sum,carry;half_adder u(a,b,sum,carry);initial begin for(integer i=0;i<4;i=i+1)begin {a,b}=i;#1;if(sum!==(a^b)||carry!==(a&b))$fatal;end $display(\"PASS\");end endmodule")
task('comb_full_adder','basic_comb','Implement module full_adder(input a, input b, input cin, output sum, output cout).',
"module tb;reg a,b,cin;wire sum,cout;full_adder u(a,b,cin,sum,cout);initial begin for(integer i=0;i<8;i=i+1)begin {a,b,cin}=i;#1;if({cout,sum}!==(a+b+cin))$fatal;end $display(\"PASS\");end endmodule")
task('comb_mux4','basic_comb','Write module mux4to1(input [7:0] a,b,c,d, input [1:0] sel, output [7:0] y).',
"module tb;reg[7:0]a,b,c,d;reg[1:0]sel;wire[7:0]y;mux4to1 u(a,b,c,d,sel,y);initial begin a=8'h11;b=8'h22;c=8'h33;d=8'h44;sel=0;#1;if(y!==a)$fatal;sel=1;#1;if(y!==b)$fatal;sel=2;#1;if(y!==c)$fatal;sel=3;#1;if(y!==d)$fatal;$display(\"PASS\");end endmodule")
task('comb_decoder','basic_comb','Write module decoder_3to8(input [2:0] in, output [7:0] out). out[in] must be the only asserted bit.',
"module tb;reg[2:0]in;wire[7:0]out;decoder_3to8 u(in,out);initial begin for(integer i=0;i<8;i=i+1)begin in=i;#1;if(out!==(8'b1<<i))$fatal;end $display(\"PASS\");end endmodule")
task('comb_priority','basic_comb','Write module priority_encoder_8to3(input [7:0] in, output [2:0] out, output valid). Highest numbered set bit wins.',
"module tb;reg[7:0]in;wire[2:0]out;wire valid;priority_encoder_8to3 u(in,out,valid);initial begin in=0;#1;if(valid!==0)$fatal;in=8'b00101000;#1;if(valid!==1||out!==5)$fatal;in=8'b10000001;#1;if(out!==7)$fatal;$display(\"PASS\");end endmodule")
task('comb_comparator','basic_comb','Implement comparator_8bit(input [7:0] a, input [7:0] b, output eq, output gt, output lt).',
"module tb;reg[7:0]a,b;wire eq,gt,lt;comparator_8bit u(a,b,eq,gt,lt);initial begin a=3;b=3;#1;if(!eq||gt||lt)$fatal;a=9;b=2;#1;if(eq||!gt||lt)$fatal;a=1;b=7;#1;if(eq||gt||!lt)$fatal;$display(\"PASS\");end endmodule")
task('comb_parity_explicit','basic_comb','Write module parity_generator(input [7:0] data, output even_parity, output odd_parity). even_parity is XOR reduction; odd_parity is inverse XOR.',
"module tb;reg[7:0]data;wire even_parity,odd_parity;parity_generator u(data,even_parity,odd_parity);initial begin data=8'h00;#1;if(even_parity!==0||odd_parity!==1)$fatal;data=8'h07;#1;if(even_parity!==1||odd_parity!==0)$fatal;data=8'h03;#1;if(even_parity!==0||odd_parity!==1)$fatal;$display(\"PASS\");end endmodule")

# --- Arithmetic / flags ---
task('arith_alu4','arith','Implement alu_4bit(input [3:0] a, input [3:0] b, input [2:0] op, output reg [3:0] y, output zero). op 0 add, 1 subtract, 2 and, 3 or, 4 xor.',
"module tb;reg[3:0]a,b;reg[2:0]op;wire[3:0]y;wire zero;alu_4bit u(a,b,op,y,zero);initial begin a=2;b=2;op=1;#1;if(y!==0||zero!==1)$fatal;op=0;#1;if(y!==4||zero!==0)$fatal;op=2;#1;if(y!==(a&b))$fatal;op=3;#1;if(y!==(a|b))$fatal;op=4;#1;if(y!==(a^b))$fatal;$display(\"PASS\");end endmodule")
task('arith_satadd','arith','Write module sat_adder8(input signed [7:0] a, input signed [7:0] b, output reg signed [7:0] sum). Saturate signed result to -128..127.',
"module tb;reg signed[7:0]a,b;wire signed[7:0]sum;sat_adder8 u(a,b,sum);initial begin a=100;b=50;#1;if(sum!==127)$fatal;a=-100;b=-50;#1;if(sum!==-128)$fatal;a=10;b=-3;#1;if(sum!==7)$fatal;$display(\"PASS\");end endmodule")
task('arith_abs','arith','Implement abs_value_8bit(input signed [7:0] in, output [7:0] out).',
"module tb;reg signed[7:0]in;wire[7:0]out;abs_value_8bit u(in,out);initial begin in=8'sd5;#1;if(out!==8'd5)$fatal;in=-8'sd5;#1;if(out!==8'd5)$fatal;$display(\"PASS\");end endmodule")

# --- Shift / rotate / bit manipulation ---
task('bit_rotate_left','bit_manip','Write module rotate_left_8(input [7:0] data, input [2:0] shamt, output [7:0] out). Rotate left by shamt.',
"module tb;reg[7:0]data;reg[2:0]shamt;wire[7:0]out;rotate_left_8 u(data,shamt,out);initial begin data=8'b10000001;shamt=0;#1;if(out!==8'b10000001)$fatal;shamt=1;#1;if(out!==8'b00000011)$fatal;shamt=4;#1;if(out!==8'b00011000)$fatal;$display(\"PASS\");end endmodule")
task('bit_barrel','bit_manip','Implement barrel_shifter_8bit(input [7:0] data, input [2:0] shamt, input dir, output [7:0] out). dir=0 shift left, dir=1 shift right logical.',
"module tb;reg[7:0]data;reg[2:0]shamt;reg dir;wire[7:0]out;barrel_shifter_8bit u(data,shamt,dir,out);initial begin data=8'b00010001;shamt=2;dir=0;#1;if(out!==8'b01000100)$fatal;dir=1;#1;if(out!==8'b00000100)$fatal;$display(\"PASS\");end endmodule")
task('bit_gray','bit_manip','Write module binary_to_gray(input [3:0] bin, output [3:0] gray).',
"module tb;reg[3:0]bin;wire[3:0]gray;binary_to_gray u(bin,gray);initial begin for(integer i=0;i<16;i=i+1)begin bin=i;#1;if(gray!==(i^(i>>1)))$fatal;end $display(\"PASS\");end endmodule")
task('bit_lod','bit_manip','Implement leading_one_detector(input [7:0] in, output reg [2:0] pos, output valid). pos is highest set bit.',
"module tb;reg[7:0]in;wire[2:0]pos;wire valid;leading_one_detector u(in,pos,valid);initial begin in=0;#1;if(valid)$fatal;in=8'b00010000;#1;if(!valid||pos!==4)$fatal;in=8'b10100000;#1;if(pos!==7)$fatal;$display(\"PASS\");end endmodule")

# --- Sequential/reset ---
task('seq_dff_sync','sequential','Write module dff_sync_reset(input clk, input rst, input d, output reg q). Reset is synchronous active-high and only updates on posedge clk.',
"module tb;reg clk=0,rst=0,d=1;wire q;dff_sync_reset u(clk,rst,d,q);always #1 clk=~clk;initial begin @(negedge clk);d=1;rst=0;@(posedge clk);#0.5;if(q!==1)$fatal;@(negedge clk);rst=1;#0.2;if(q!==1)$fatal;@(posedge clk);#0.5;if(q!==0)$fatal;rst=0;d=1;@(posedge clk);#0.5;if(q!==1)$fatal;$display(\"PASS\");$finish;end endmodule")
task('seq_dff_async','sequential','Write module dff_async_reset(input clk, input rst, input d, output reg q). Reset is asynchronous active-high.',
"module tb;reg clk=0,rst=0,d=1;wire q;dff_async_reset u(clk,rst,d,q);always #1 clk=~clk;initial begin @(posedge clk);#0.5;if(q!==1)$fatal;rst=1;#0.2;if(q!==0)$fatal;rst=0;d=1;@(posedge clk);#0.5;if(q!==1)$fatal;$display(\"PASS\");$finish;end endmodule")
task('seq_counter','sequential','Write module counter_8bit_up(input clk, input rst, input en, output reg [7:0] count). Synchronous reset; increment when en.',
"module tb;reg clk=0,rst=1,en=0;wire[7:0]count;counter_8bit_up u(clk,rst,en,count);always #1 clk=~clk;initial begin @(negedge clk);rst=1;en=0;@(posedge clk);#0.5;if(count!==0)$fatal;@(negedge clk);rst=0;en=1;repeat(5)@(posedge clk);#0.5;if(count!==5)$fatal;@(negedge clk);en=0;repeat(2)@(posedge clk);#0.5;if(count!==5)$fatal;$display(\"PASS\");$finish;end endmodule")
task('seq_shiftreg','sequential','Implement shift_register_8(input clk, input rst, input serial_in, output reg [7:0] data). Shift left and insert serial_in at bit0.',
"module tb;reg clk=0,rst=1,serial_in=0;wire[7:0]data;shift_register_8 u(clk,rst,serial_in,data);always #1 clk=~clk;task send(input b);begin @(negedge clk);serial_in=b;@(posedge clk);#0.5;end endtask initial begin @(posedge clk);rst=0;send(1);send(0);send(1);if(data[2:0]!==3'b101)$fatal;$display(\"PASS\");$finish;end endmodule")
task('seq_edge','sequential','Write module rising_edge_detector(input clk, input rst, input signal_in, output reg pulse). pulse high one clock when signal_in rises.',
"module tb;reg clk=0,rst=1,signal_in=0;wire pulse;rising_edge_detector u(clk,rst,signal_in,pulse);always #1 clk=~clk;initial begin @(negedge clk);rst=0;signal_in=0;@(posedge clk);@(negedge clk);signal_in=1;@(posedge clk);#0.5;if(!pulse)$fatal;@(posedge clk);#0.5;if(pulse)$fatal;$display(\"PASS\");$finish;end endmodule")

# --- FSM ---
task('fsm_1011_overlap','fsm','Write module sequence_detector_1011(input clk, input rst, input bit_in, output reg detected). Detect overlapping pattern 1011; assert when final bit arrives.',
"module tb;reg clk=0,rst=1,bit_in=0;wire detected;sequence_detector_1011 u(clk,rst,bit_in,detected);always #1 clk=~clk;task send(input b);begin @(negedge clk);bit_in=b;@(posedge clk);#0.5;end endtask initial begin @(posedge clk);rst=0;send(1);send(0);send(1);send(1);if(!detected)$fatal;send(0);send(1);send(1);if(!detected)$fatal;$display(\"PASS\");$finish;end endmodule")
task('fsm_traffic','fsm','Write traffic_light_fsm(input clk, input rst, output reg [1:0] ns_light, output reg [1:0] ew_light). Cycle NS green, NS yellow, EW green, EW yellow.', None, True)
task('fsm_vending','fsm','Write vending_machine(input clk, input rst, input nickel, input dime, output reg dispense, output reg [3:0] credit). Dispense when credit reaches at least 15 cents.', None, True)

# --- Memory ---
task('mem_regfile','memory','Implement register_file_32x32(input clk, input we, input [4:0] rd_addr1, input [4:0] rd_addr2, input [4:0] wr_addr, input [31:0] wr_data, output [31:0] rd_data1, output [31:0] rd_data2). Two async read ports, sync write, register 0 always zero.',
"module tb;reg clk=0,we=0;reg[4:0]rd_addr1=0,rd_addr2=0,wr_addr=0;reg[31:0]wr_data=0;wire[31:0]rd_data1,rd_data2;register_file_32x32 u(clk,we,rd_addr1,rd_addr2,wr_addr,wr_data,rd_data1,rd_data2);always #1 clk=~clk;initial begin @(negedge clk);we=1;wr_addr=0;wr_data=32'hffff;@(posedge clk);#0.5;rd_addr1=0;#1;if(rd_data1!==0)$fatal;@(negedge clk);wr_addr=9;wr_data=32'h12345678;we=1;@(posedge clk);#0.5;we=0;rd_addr2=9;#1;if(rd_data2!==32'h12345678)$fatal;$display(\"PASS\");$finish;end endmodule")
task('mem_ram16','memory','Write module ram_16x8(input clk, input we, input [3:0] addr, input [7:0] din, output [7:0] dout). Synchronous write, asynchronous read.',
"module tb;reg clk=0,we=0;reg[3:0]addr=0;reg[7:0]din=0;wire[7:0]dout;ram_16x8 u(clk,we,addr,din,dout);always #1 clk=~clk;initial begin @(negedge clk);addr=4'd3;din=8'ha5;we=1;@(posedge clk);#0.5;we=0;#1;if(dout!==8'ha5)$fatal;$display(\"PASS\");$finish;end endmodule")
task('mem_dp_ram','memory','Write dp_ram_32x8(input clk, input we_a, input [4:0] addr_a, input [7:0] din_a, output [7:0] dout_a, input we_b, input [4:0] addr_b, input [7:0] din_b, output [7:0] dout_b).', None, True)

# --- Protocol/compile-heavy RTLLM-ish ---
task('rtl_spi','compile_rtl','Write complete synthesizable Verilog for spi_master(input clk, input rst, input start, input [7:0] data_in, output reg sclk, output reg mosi, output reg busy, output reg done).', None, True)
task('rtl_uart','compile_rtl','Write complete synthesizable Verilog for uart_tx(input clk, input rst, input start, input [7:0] data, output reg tx, output reg busy, output reg done).', None, True)
task('rtl_pwm','compile_rtl','Write pwm_generator(input clk, input rst, input [7:0] duty, output reg pwm). Use an 8-bit counter and pwm high when counter < duty.', None, True)
task('rtl_crc8','compile_rtl','Write crc8_generator(input clk, input rst, input en, input [7:0] data_in, output reg [7:0] crc). Use polynomial 0x07 update when en.', None, True)
task('rtl_rr_arbiter','compile_rtl','Write rr_arbiter4(input clk, input rst, input [3:0] req, output reg [3:0] grant). Grant one request using round-robin priority.', None, True)

def extract_code(text):
    for m in re.finditer(r"\[BEGIN\](.*?)\[DONE\]", text, re.S|re.I):
        span=m.group(1)
        if 'module ' in span and 'endmodule' in span:
            text=span; break
    else:
        m=re.search(r"```(?:verilog|systemverilog)?\s*(.*?)```", text, re.S|re.I)
        if m: text=m.group(1)
    i=text.find('module ')
    if i>=0:
        text=text[i:]
        j=text.rfind('endmodule')
        if j>=0: text=text[:j+9]
    return text.strip()

def run_iverilog(code, tb):
    with tempfile.TemporaryDirectory() as d:
        dut=Path(d)/'dut.v'; dut.write_text(code)
        files=[str(dut)]
        if tb:
            tbp=Path(d)/'tb.v'; tbp.write_text(tb); files.append(str(tbp))
        out=Path(d)/'sim.out'
        c=subprocess.run(['iverilog','-g2012','-o',str(out)]+files,capture_output=True,text=True,timeout=30)
        if c.returncode!=0: return False, False, (c.stderr+c.stdout)[-3000:]
        if not tb: return True, None, ''
        v=subprocess.run(['vvp',str(out)],capture_output=True,text=True,timeout=30)
        txt=v.stdout+v.stderr
        return True, (v.returncode==0 and 'PASS' in txt and 'FAIL' not in txt), txt[-3000:]

def split_hf_subfolder(path):
    """Allow adapter refs like owner/repo/subfolder in addition to local paths."""
    if path and not Path(path).exists() and path.count('/') >= 2:
        parts = path.split('/')
        return '/'.join(parts[:2]), '/'.join(parts[2:])
    return path, None


def load_model(base, adapter):
    import torch
    from peft import PeftModel
    from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText, AutoTokenizer, BitsAndBytesConfig
    adapter_repo, adapter_subfolder = split_hf_subfolder(adapter) if adapter else (None, None)
    # Always use the base tokenizer. Some adapter subfolders may not carry a
    # complete tokenizer/config set, and tokenizer must match base model anyway.
    tok=AutoTokenizer.from_pretrained(base, trust_remote_code=True)
    tok.pad_token=tok.eos_token
    dtype=torch.bfloat16 if 'Qwen3.5' in base or 'Qwen3' in base else torch.float16
    bnb=BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_quant_type='nf4',bnb_4bit_compute_dtype=dtype,bnb_4bit_use_double_quant=True)
    cfg=AutoConfig.from_pretrained(base, trust_remote_code=True)
    model_type=getattr(cfg,'model_type','')
    arch=' '.join(getattr(cfg,'architectures',[]) or [])
    loader=AutoModelForImageTextToText if (model_type in {'qwen3_5'} or 'ForConditionalGeneration' in arch) else AutoModelForCausalLM
    model=loader.from_pretrained(base,quantization_config=bnb,device_map='auto',trust_remote_code=True)
    if adapter_repo:
        if adapter_subfolder:
            model=PeftModel.from_pretrained(model, adapter_repo, subfolder=adapter_subfolder)
        else:
            model=PeftModel.from_pretrained(model, adapter_repo)
    model.eval(); return model,tok

def gen(model,tok,prompt,max_new_tokens,temp,top_p,seed):
    import torch
    if seed is not None:
        torch.manual_seed(seed); random.seed(seed)
    msg=[{'role':'system','content':SYSTEM},{'role':'user','content':prompt}]
    text=tok.apply_chat_template(msg,tokenize=False,add_generation_prompt=True)
    inp=tok(text,return_tensors='pt').to(model.device)
    sample=temp>0
    with torch.no_grad():
        out=model.generate(**inp,max_new_tokens=max_new_tokens,do_sample=sample,temperature=temp if sample else None,top_p=top_p if sample else None,pad_token_id=tok.eos_token_id)
    return tok.decode(out[0][inp['input_ids'].shape[-1]:],skip_special_tokens=True)

def pass_at_k(n,c,k):
    if n-c < k: return 1.0
    prod=1.0
    for i in range(k): prod *= (n-c-i)/(n-i)
    return 1.0-prod

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument('--suite',default='full',choices=['full','smoke'])
    ap.add_argument('--base',default=BASE)
    ap.add_argument('--adapter')
    ap.add_argument('--model-name',default='model')
    ap.add_argument('--out',default='internal_rtl_diagnostic_results/model')
    ap.add_argument('--k',type=int,default=1,help='samples per task')
    ap.add_argument('--temperature',type=float,default=0.0)
    ap.add_argument('--top-p',type=float,default=0.95)
    ap.add_argument('--max-new-tokens',type=int,default=700)
    ap.add_argument('--list-suites',action='store_true')
    args=ap.parse_args()
    if args.list_suites:
        cats={}
        for t in TASKS: cats[t['category']]=cats.get(t['category'],0)+1
        print(json.dumps({'full':len(TASKS),'smoke':8,'categories':cats},indent=2)); return
    tasks=TASKS[:8] if args.suite=='smoke' else TASKS
    od=Path(args.out); od.mkdir(parents=True,exist_ok=True)
    model,tok=load_model(args.base,args.adapter)
    records=[]
    for t in tasks:
        successes=0; compile_successes=0; samples=[]
        for s in range(args.k):
            t0=time.time(); raw=gen(model,tok,t['prompt'],args.max_new_tokens,args.temperature,args.top_p,seed=1000+s if args.temperature>0 else None)
            code=extract_code(raw); comp,func,err=run_iverilog(code,t['testbench'])
            ok=comp and (func is True if t['testbench'] else True)
            successes += int(ok); compile_successes += int(comp)
            sample={'sample':s,'compile':comp,'functional':func,'ok':ok,'sec':round(time.time()-t0,2),'code':code,'err':err}
            samples.append(sample); (od/f"{t['id']}_s{s}.v").write_text(code)
            print(json.dumps({'id':t['id'],'sample':s,'compile':comp,'functional':func,'ok':ok,'sec':sample['sec']}),flush=True)
        rec={**t,'samples':samples,'pass':successes>0,'compile_any':compile_successes>0,'successes':successes,'compile_successes':compile_successes}
        records.append(rec)
    total=len(records); func_total=sum(1 for r in records if r['testbench'])
    summary={'model':args.model_name,'adapter':args.adapter,'suite':args.suite,'tasks':total,'k':args.k,
             'compile_any_pass':sum(r['compile_any'] for r in records),'compile_any_pct':100*sum(r['compile_any'] for r in records)/total,
             'task_pass':sum(r['pass'] for r in records),'task_pass_pct':100*sum(r['pass'] for r in records)/total,
             'functional_total':func_total,'functional_pass':sum(r['pass'] for r in records if r['testbench']),
             'functional_pct':100*sum(r['pass'] for r in records if r['testbench'])/func_total if func_total else None}
    if args.k>1:
        summary['pass_at_1_est']=sum(pass_at_k(args.k,r['successes'],1) for r in records)/total
        summary[f'pass_at_{args.k}_est']=sum(pass_at_k(args.k,r['successes'],args.k) for r in records)/total
    bycat={}
    for r in records:
        c=r['category']; bycat.setdefault(c,{'total':0,'pass':0,'compile':0}); bycat[c]['total']+=1; bycat[c]['pass']+=int(r['pass']); bycat[c]['compile']+=int(r['compile_any'])
    for c,v in bycat.items(): v['pass_pct']=100*v['pass']/v['total']; v['compile_pct']=100*v['compile']/v['total']
    summary['by_category']=bycat
    (od/'results.jsonl').write_text(''.join(json.dumps(r)+'\n' for r in records))
    (od/'summary.json').write_text(json.dumps(summary,indent=2))
    print(json.dumps(summary,indent=2))
if __name__=='__main__': main()