"""The generated RTL carries the calibrated constants and computes the decision.

The constant and structural checks run everywhere. The simulation checks run
when Icarus Verilog is available, either on PATH or via the IVERILOG and VVP
environment variables.
"""
import os
import re
import shutil
import subprocess

import pytest

from conftest import REPO, load

RTL = REPO / 'rtl'
N_DIMS = 40
N_POS = 20
N_VECTORS = 512


def tool(name, env_var):
    return os.environ.get(env_var) or shutil.which(name)


IVERILOG = tool('iverilog', 'IVERILOG')
VVP = tool('vvp', 'VVP')
needs_sim = pytest.mark.skipif(not (IVERILOG and VVP),
                               reason='Icarus Verilog not found on PATH or in IVERILOG/VVP')


@pytest.fixture(scope='module')
def calibration():
    return load('per_dim_thresholds.json')


def source(name: str) -> str:
    return (RTL / f'{name}.v').read_text(encoding='utf-8')


# ---------------- constants and structure ----------------

def test_baked_per_dim_thresholds_match_the_calibration(calibration):
    text = source('popcount_folded')
    baked = {int(i): int(v) for i, v in re.findall(r'T(\d\d) =\s*(-?\d+)', text)}
    assert len(baked) == N_DIMS
    for e in calibration['per_dim_thresholds']:
        assert baked[e['dim_index_in_40']] == e['threshold_int8']


def test_baked_final_threshold_matches_the_calibration(calibration):
    text = source('popcount_folded')
    assert int(re.search(r'FINAL_T = (-?\d+);', text).group(1)) == \
        calibration['popcount']['final_threshold']


def test_sum_folded_threshold_is_the_quantized_threshold(classifier, calibration):
    text = source('sum_folded')
    baked = int(re.search(r"FINAL_T = 16'sd(-?\d+);", text).group(1))
    assert baked == round(classifier['threshold'] * calibration['quant_scale'])


@pytest.mark.parametrize('name,threshold', [('popcount', 't{:02d}'),
                                            ('popcount_folded', 'T{:02d}')])
def test_each_channel_is_compared_against_its_own_threshold(name, threshold):
    text = source(name)
    pairs = re.findall(r'\(f(\d\d) > (\w+)\)', text)
    assert len(pairs) == N_DIMS
    for feature, thr in pairs:
        assert thr == threshold.format(int(feature))
    assert [int(f) for f, _ in pairs] == list(range(N_DIMS))


@pytest.mark.parametrize('name', ['popcount', 'popcount_folded'])
def test_no_bit_select_on_a_vector(name):
    """Bit-select lowering is unreliable at index 0 in the synthesis backend.

    The comparisons are summed inline instead, so nothing here depends on it.
    """
    text = '\n'.join(l for l in source(name).splitlines()
                     if not l.strip().startswith('//'))
    assert not re.search(r'\b(pos|neg)_bits\b', text)


@pytest.mark.parametrize('name', ['sum', 'sum_folded', 'popcount', 'popcount_folded'])
def test_every_feature_port_is_declared_once(name):
    declared = re.findall(r'\bf(\d\d)\b(?=[,)\s])', source(name).split(');')[0])
    assert sorted(set(declared)) == [f'{i:02d}' for i in range(N_DIMS)]


# ---------------- simulation ----------------

def vectors(n: int, seed: int = 0):
    """Deterministic signed-INT8 feature vectors, uniform over the input range."""
    import random
    rng = random.Random(seed)
    return [[rng.randint(-128, 127) for _ in range(N_DIMS)] for _ in range(n)]


def popcount_boundary_vectors(thresholds, k: int, n: int, seed: int = 0):
    """Vectors placing the count difference within two of K, on the per-dim thresholds.

    An `off` channel sits at exactly its threshold, so `>` must reject it.
    """
    import random
    rng = random.Random(seed)
    out = []
    while len(out) < n:
        diff = k + rng.randint(-2, 2)
        n_neg = rng.randint(0, max(0, N_POS - abs(diff)))
        n_pos = diff + n_neg
        if not (0 <= n_pos <= N_POS and 0 <= n_neg <= N_POS):
            continue
        pos_on = [True] * n_pos + [False] * (N_POS - n_pos)
        neg_on = [True] * n_neg + [False] * (N_POS - n_neg)
        rng.shuffle(pos_on)
        rng.shuffle(neg_on)
        on = pos_on + neg_on
        vec = [max(-128, min(127, t + 1)) if on[i] else max(-128, min(127, t))
               for i, t in enumerate(thresholds)]
        out.append(vec)
    return out


def additive_boundary_vectors(threshold: int, n: int, seed: int = 0):
    """Vectors whose signed sum lands within a few counts of the comparator threshold."""
    import random
    rng = random.Random(seed)
    out = []
    while len(out) < n:
        vec = [rng.randint(-20, 20) for _ in range(N_DIMS)]
        target = threshold + rng.randint(-4, 4)
        # Solve f00 so that sum(pos) - sum(neg) hits the target exactly.
        vec[0] = target - (sum(vec[1:N_POS]) - sum(vec[N_POS:]))
        if -128 <= vec[0] <= 127:
            out.append(vec)
    return out


def additive_reference(vec, threshold):
    return (sum(vec[:N_POS]) - sum(vec[N_POS:])) > threshold


def popcount_reference(vec, thresholds, k):
    pos = sum(1 for i in range(N_POS) if vec[i] > thresholds[i])
    neg = sum(1 for i in range(N_POS, N_DIMS) if vec[i] > thresholds[i])
    return (pos - neg) > k


def literal(value: int, width: int) -> str:
    return f"-{width}'sd{-value}" if value < 0 else f"{width}'sd{value}"


def make_testbench(module: str, extra_ports: str, n: int) -> str:
    features = ', '.join(f'f{i:02d}' for i in range(N_DIMS))
    connections = ',\n        '.join(f'.f{i:02d}(f{i:02d})' for i in range(N_DIMS))
    slices = '\n      '.join(
        f'f{i:02d} = vecs[i][{319 - 8 * i}:{312 - 8 * i}];' for i in range(N_DIMS))
    return f'''`timescale 1ns/1ps
module tb;
  reg [319:0] vecs [0:{n - 1}];
  reg signed [7:0] {features};
  wire out;
  integer i;
  {module} dut (
        {connections},{extra_ports}
        .person_present(out));
  initial begin
    $readmemh("vectors.hex", vecs);
    for (i = 0; i < {n}; i = i + 1) begin
      {slices}
      #1;
      $display("%b", out);
    end
    $finish;
  end
endmodule
'''


def run_sim(tmp_path, rtl_name, module, extra_ports, vecs):
    tmp_path.mkdir(parents=True, exist_ok=True)
    hexfile = tmp_path / 'vectors.hex'
    hexfile.write_text('\n'.join(
        ''.join(f'{v & 0xFF:02x}' for v in vec) for vec in vecs) + '\n')
    (tmp_path / 'tb.v').write_text(make_testbench(module, extra_ports, len(vecs)))
    subprocess.run([IVERILOG, '-g2005', '-o', 'tb.vvp',
                    str(RTL / f'{rtl_name}.v'), 'tb.v'],
                   cwd=tmp_path, check=True, capture_output=True)
    out = subprocess.run([VVP, 'tb.vvp'], cwd=tmp_path, check=True,
                         capture_output=True, text=True).stdout
    bits = [line.strip() for line in out.splitlines() if line.strip() in ('0', '1')]
    assert len(bits) == len(vecs), f'{module}: {len(bits)} results for {len(vecs)} vectors'
    return [b == '1' for b in bits]


@needs_sim
def test_sum_matches_the_additive_reference(tmp_path, classifier, calibration):
    thr = round(classifier['threshold'] * calibration['quant_scale'])
    vecs = vectors(N_VECTORS) + additive_boundary_vectors(thr, N_VECTORS)
    got = run_sim(tmp_path, 'sum', 'person_classifier_1p',
                  f'\n        .threshold({literal(thr, 16)}),', vecs)
    assert got == [additive_reference(v, thr) for v in vecs]


@needs_sim
def test_sum_folded_matches_the_additive_reference(tmp_path, classifier, calibration):
    thr = round(classifier['threshold'] * calibration['quant_scale'])
    vecs = vectors(N_VECTORS) + additive_boundary_vectors(thr, N_VECTORS)
    got = run_sim(tmp_path, 'sum_folded', 'person_classifier_sum_folded', '', vecs)
    assert got == [additive_reference(v, thr) for v in vecs]


@needs_sim
def test_popcount_matches_the_popcount_reference(tmp_path, calibration):
    thresholds = [e['threshold_int8'] for e in calibration['per_dim_thresholds']]
    k = calibration['popcount']['final_threshold']
    ports = ''.join(f'\n        .t{i:02d}({literal(t, 8)}),'
                    for i, t in enumerate(thresholds))
    ports += f'\n        .final_threshold({literal(k, 6)}),'
    vecs = vectors(N_VECTORS) + popcount_boundary_vectors(thresholds, k, N_VECTORS)
    got = run_sim(tmp_path, 'popcount', 'person_classifier_popcount', ports, vecs)
    assert got == [popcount_reference(v, thresholds, k) for v in vecs]


@needs_sim
def test_popcount_folded_matches_the_popcount_reference(tmp_path, calibration):
    thresholds = [e['threshold_int8'] for e in calibration['per_dim_thresholds']]
    k = calibration['popcount']['final_threshold']
    vecs = vectors(N_VECTORS) + popcount_boundary_vectors(thresholds, k, N_VECTORS)
    got = run_sim(tmp_path, 'popcount_folded', 'person_classifier_popcount_folded',
                  '', vecs)
    assert got == [popcount_reference(v, thresholds, k) for v in vecs]


@needs_sim
def test_folding_a_threshold_does_not_change_the_decision(tmp_path, classifier,
                                                          calibration):
    """Each runtime-threshold module agrees with its baked counterpart."""
    thr = round(classifier['threshold'] * calibration['quant_scale'])
    vecs = vectors(N_VECTORS, seed=1) + additive_boundary_vectors(thr, N_VECTORS, seed=1)
    runtime = run_sim(tmp_path / 'a', 'sum', 'person_classifier_1p',
                      f'\n        .threshold({literal(thr, 16)}),', vecs)
    baked = run_sim(tmp_path / 'b', 'sum_folded', 'person_classifier_sum_folded', '', vecs)
    assert runtime == baked