"""Generate synthetic quantum syndrome data for training."""
import numpy as np
import json
from pathlib import Path


def generate_surface_code_syndrome(distance=5, error_rate=0.05, num_samples=1000):
    """
    Generate syndrome data for a dxd surface code.
    
    Returns:
        syndromes: (num_samples, distance, distance, 2) array
            Channel 0: X syndrome (star stabilizers)
            Channel 1: Z syndrome (plaquette stabilizers)
        errors: (num_samples, distance, distance, 2) array
            Ground truth error locations
    """
    syndromes = np.zeros((num_samples, distance, distance, 2), dtype=np.float32)
    errors = np.zeros((num_samples, distance, distance, 2), dtype=np.float32)
    
    for s in range(num_samples):
        # Random X and Z errors on data qubits
        x_errors = np.random.random((distance, distance)) < error_rate
        z_errors = np.random.random((distance, distance)) < error_rate
        
        errors[s, :, :, 0] = x_errors.astype(np.float32)
        errors[s, :, :, 1] = z_errors.astype(np.float32)
        
        # Compute X syndrome (star stabilizers measure Z errors)
        for i in range(distance):
            for j in range(distance):
                # Star stabilizer at (i,j) checks Z errors on 4 neighboring qubits
                z_synd = 0
                for di, dj in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                    ni, nj = i + di, j + dj
                    if 0 <= ni < distance and 0 <= nj < distance:
                        z_synd ^= int(z_errors[ni, nj])
                syndromes[s, i, j, 0] = z_synd
                
        # Compute Z syndrome (plaquette stabilizers measure X errors)
        for i in range(distance - 1):
            for j in range(distance - 1):
                # Plaquette stabilizer at (i,j) checks X errors on 4 qubits
                x_synd = 0
                for di, dj in [(0, 0), (0, 1), (1, 0), (1, 1)]:
                    x_synd ^= int(x_errors[i + di, j + dj])
                syndromes[s, i, j, 1] = x_synd
    
    return syndromes, errors


def generate_dataset(distances=[5, 7, 9], error_rates=[0.01, 0.03, 0.05, 0.08], 
                     samples_per_config=5000, output_dir="data"):
    """Generate full training dataset."""
    Path(output_dir).mkdir(exist_ok=True)
    
    all_syndromes = []
    all_errors = []
    all_metadata = []
    
    for d in distances:
        for p in error_rates:
            print(f"Generating d={d}, p={p}...")
            synd, errs = generate_surface_code_syndrome(d, p, samples_per_config)
            all_syndromes.append(synd)
            all_errors.append(errs)
            all_metadata.extend([{"distance": d, "error_rate": p}] * samples_per_config)
    
    # Pad to max distance
    max_d = max(distances)
    padded_syndromes = []
    padded_errors = []
    
    for synd, errs in zip(all_syndromes, all_errors):
        d = synd.shape[1]
        if d < max_d:
            # Pad with zeros
            pad = ((0, 0), (0, max_d - d), (0, max_d - d), (0, 0))
            synd = np.pad(synd, pad, mode='constant')
            errs = np.pad(errs, pad, mode='constant')
        padded_syndromes.append(synd)
        padded_errors.append(errs)
    
    syndromes = np.concatenate(padded_syndromes, axis=0)
    errors = np.concatenate(padded_errors, axis=0)
    
    # Save
    np.save(f"{output_dir}/syndromes.npy", syndromes)
    np.save(f"{output_dir}/errors.npy", errors)
    with open(f"{output_dir}/metadata.json", "w") as f:
        json.dump(all_metadata, f)
    
    print(f"Dataset saved: {len(syndromes)} samples, shape={syndromes.shape}")
    return syndromes, errors, all_metadata


if __name__ == "__main__":
    generate_dataset()