import numpy as np
import cv2
import random
import os

def generate_circuit_image(size=512):
    canvas = np.zeros((size, size), dtype=np.uint8)
    
    if random.random() > 0.7:
        words = ["DREAM", "CORE", "SYNC", "DATA", "VAE", "BUS", "GATE", "NODE", "LINK"]
        word = random.choice(words)
        font = cv2.FONT_HERSHEY_SIMPLEX
        scale = random.uniform(1.0, 2.5)
        thickness = random.randint(2, 4)
        pos = (random.randint(50, size-200), random.randint(50, size-50))
        cv2.putText(canvas, word, pos, font, scale, 255, thickness, cv2.LINE_AA)

    num_paths = random.randint(15, 25)
    
    for _ in range(num_paths):
        start_x, start_y = random.randint(20, size-20), random.randint(20, size-20)
        thickness = random.choice([1, 2, 3, 5])
        
        curr_x, curr_y = start_x, start_y
        points = [(curr_x, curr_y)]
        
        num_segments = random.randint(3, 8)
        for _ in range(num_segments):
            direction = random.choice(['h', 'v', 'd'])
            length = random.randint(20, 100)
            
            if direction == 'h':
                curr_x = np.clip(curr_x + random.choice([-1, 1]) * length, 20, size-20)
            elif direction == 'v':
                curr_y = np.clip(curr_y + random.choice([-1, 1]) * length, 20, size-20)
            else:
                d = random.choice([-1, 1]) * length
                curr_x = np.clip(curr_x + d, 20, size-20)
                curr_y = np.clip(curr_y + d, 20, size-20)
            
            points.append((int(curr_x), int(curr_y)))
            
        for i in range(len(points) - 1):
            cv2.line(canvas, points[i], points[i+1], 255, thickness)
            
        cv2.circle(canvas, points[0], thickness + random.randint(2, 5), 255, -1)
        cv2.circle(canvas, points[-1], thickness + random.randint(2, 5), 255, -1)
        
        for p in points[1:-1]:
            cv2.circle(canvas, p, thickness + 1, 255, -1)

    return canvas

def generate_batch(batch_size, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    for i in range(batch_size):
        img = generate_circuit_image()
        cv2.imwrite(os.path.join(output_dir, f'circuit_{i:04d}.png'), img)
        if (i+1) % 50 == 0:
            print(f"GEN_PROG: {i+1}/{batch_size}")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--count", type=int, default=100)
    parser.add_argument("--out", type=str, default="data/training")
    args = parser.parse_args()
    
    print(f"Generating {args.count} CIRCUIT samples to {args.out}...")
    generate_batch(args.count, args.out)
    print("Done.")
