import json

def rotate_90(position):
    """
    Rotate (pinA, pinB) by 90 degrees around the center of a 5x5 board.
    Formula: (pinA, pinB) -> (pinB, 4 - pinA)
    """
    pinA, pinB = position
    return (pinB, 4 - pinA)

def rotate_all_rotations(pins):
    """
    Generate all unique rotations (90, 180, 270 degrees) of the given pin positions.
    Returns a list of rotated pin sets.
    """
    rotated_pins = []
    current_pins = pins.copy()
    for _ in range(3):  # Rotate 3 times (90, 180, 270 degrees)
        current_pins = [rotate_90(pin) for pin in current_pins]
        rotated_pins.append(current_pins.copy())
    return rotated_pins

def pins_to_tuple(entry):
    # Convert each pin's coordinates to tuples for immutability
    return (
        tuple(entry['pinA']),
        tuple(entry['pinB']),
        tuple(entry['pinC'])
    )

def main():
    # Load existing data
    with open('data.json', 'r') as f:
        data = json.load(f)
    
    # Initialize a set to track existing forbidden pin sets to avoid duplicates
    existing_pins_set = set()
    
    # Populate the existing_pins_set with the original data
    for entry in data:
        pins_tuple = pins_to_tuple(entry)
        existing_pins_set.add(pins_tuple)
    
    # List to hold new rotated entries
    new_entries = []
    
    # Iterate over the existing data and generate rotated entries
    for entry in data:
        pinA = tuple(entry['pinA'])
        pinB = tuple(entry['pinB'])
        pinC = tuple(entry['pinC'])
        
        base_pins = [pinA, pinB, pinC]
        
        # Generate all rotated versions
        rotated_versions = rotate_all_rotations(base_pins)
        
        for rotated in rotated_versions:
            rotated_pinA, rotated_pinB, rotated_pinC = rotated
            rotated_entry = {
                'pinA': list(rotated_pinA),
                'pinB': list(rotated_pinB),
                'pinC': list(rotated_pinC),
                'solution_count': entry['solution_count']
            }
            rotated_pins_tuple = (
                tuple(rotated_entry['pinA']),
                tuple(rotated_entry['pinB']),
                tuple(rotated_entry['pinC'])
            )
            if rotated_pins_tuple not in existing_pins_set:
                new_entries.append(rotated_entry)
                existing_pins_set.add(rotated_pins_tuple)
    
    print(f"Added {len(new_entries)} rotated entries.")
    
    # Combine original data with new rotated entries
    extended_data = data + new_entries
    
    # Define the arrays for pinA, pinB, and pinC containing tuples of coordinates
    pinA = [(0, 0), (1, 0), (0, 1), (1, 1), (2, 1), (3, 1), (1, 2), (1, 3)]  # x pins
    pinB = [(3, 0), (4, 0), (4, 1), (4, 2), (4, 3), (3, 4), (4, 4)]          # y pins
    pinC = [(2, 2), (3, 2), (0, 3), (2, 3), (0, 4), (1, 4), (2, 4)]          # z pins

    # Define the array for forbidden positions
    forbidden_positions = [(2, 0), (0, 2), (3, 3)]  # o positions

    # Remove forbidden positions from the extended data
    filtered_data = []
    for entry in extended_data:
        pinA, pinB, pinC = entry['pinA'], entry['pinB'], entry['pinC']
        if pinA not in forbidden_positions and pinB not in forbidden_positions and pinC not in forbidden_positions:
            filtered_data.append(entry)
    
    with open('data_extended.json', 'w') as f:
        json.dump(extended_data, f, indent=4)
    
    with open("data_filtered.json", "w") as f:
        json.dump(filtered_data, f, indent=4)
        
    # Print distribution of solutions in the filtered data
    solution_counts = [entry['solution_count'] for entry in filtered_data]
    unique_solution_counts = set(solution_counts)
    print("Distribution of solutions in the filtered data:")
    for count in unique_solution_counts:
        count_entries = [entry for entry in filtered_data if entry['solution_count'] == count]
        print(f"Solutions: {count}, Count: {len(count_entries)}")
    
if __name__ == "__main__":
    main()
