#!/usr/bin/env python3
"""
Generate golden test vectors for Sobel edge detection testbench.
Uses standard Sobel kernels and |Gx| + |Gy| approximation for gradient magnitude.
"""

import numpy as np

# Image dimensions
IMG_H, IMG_W = 64, 64

# Standard Sobel kernels
SOBEL_GX = np.array([
    [-1, 0, 1],
    [-2, 0, 2],
    [-1, 0, 1]
], dtype=np.int8)

SOBEL_GY = np.array([
    [-1, -2, -1],
    [ 0,  0,  0],
    [ 1,  2,  1]
], dtype=np.int8)


def generate_test_input():
    """Generate a 64x64 test image with gradient pattern."""
    img = np.zeros((IMG_H, IMG_W), dtype=np.uint8)
    for r in range(IMG_H):
        for c in range(IMG_W):
            # Diagonal gradient pattern - good for edge detection testing
            img[r, c] = (r * 4 + c * 2) % 256
    return img


def sobel_edge_detection(img, pad=1):
    """
    Compute Sobel edge detection with zero-padding.
    Uses |Gx| + |Gy| approximation for gradient magnitude.
    Returns 8-bit unsigned output (clamped to [0, 255]).
    """
    h, w = img.shape
    
    # Zero-pad the input
    img_padded = np.pad(img.astype(np.int32), pad, mode='constant', constant_values=0)
    
    # Output dimensions (same as input for stride=1, pad=1)
    out_h = h
    out_w = w
    output = np.zeros((out_h, out_w), dtype=np.int32)
    
    # Apply Sobel convolution
    for r in range(out_h):
        for c in range(out_w):
            # Extract 3x3 region
            region = img_padded[r:r+3, c:c+3]
            
            # Compute Gx and Gy
            gx = np.sum(region * SOBEL_GX.astype(np.int32))
            gy = np.sum(region * SOBEL_GY.astype(np.int32))
            
            # Gradient magnitude: |Gx| + |Gy|
            magnitude = abs(gx) + abs(gy)
            
            # Saturate to [0, 255]
            output[r, c] = min(magnitude, 255)
    
    return output.astype(np.uint8)


def generate_testbench_array():
    """Generate expected_output as compact Verilog array literal for testbench embedding."""
    img = generate_test_input()
    output = sobel_edge_detection(img)
    
    # Collect all values
    values = []
    for r in range(IMG_H):
        for c in range(IMG_W):
            values.append(f"8'h{output[r, c]:02X}")
    
    # Print as compact array (16 values per line)
    print("    // Golden expected outputs from generate_golden.py")
    print("    // DO NOT EDIT - regenerate with: python3 generate_golden.py --array")
    print(f"    reg [7:0] expected_output [0:{IMG_H*IMG_W-1}] = '{{")
    for i in range(0, len(values), 16):
        line_vals = values[i:i+16]
        suffix = "," if i + 16 < len(values) else ""
        print("        " + ", ".join(line_vals) + suffix)
    print("    };")


def generate_verbose():
    """Generate full verbose output for debugging."""
    img = generate_test_input()
    output = sobel_edge_detection(img)
    
    print("// ============================================")
    print("// AUTO-GENERATED TEST VECTORS - DO NOT EDIT")
    print("// Run: python3 generate_golden.py")
    print("// ============================================")
    print()
    
    print("// Sobel kernels (standard)")
    print("// Gx: [[-1,0,1], [-2,0,2], [-1,0,1]]")
    print("// Gy: [[-1,-2,-1], [0,0,0], [1,2,1]]")
    print()
    
    # Input image (first 10 values for reference)
    print(f"// Test image: {IMG_H}x{IMG_W} gradient pattern (r*4 + c*2) % 256")
    print(f"// First row: {list(img[0, :10])}...")
    print()
    
    # Expected output (first 10 values for reference)
    print(f"// Expected output: {IMG_H}x{IMG_W} edge magnitudes")
    print(f"// First row: {list(output[0, :10])}...")
    print()
    
    # Statistics
    print(f"// Output statistics:")
    print(f"//   Min: {output.min()}, Max: {output.max()}, Mean: {output.mean():.1f}")


if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "--array":
        generate_testbench_array()
    else:
        generate_verbose()