main.py

import os
import yaml
import logging
import concurrent.futures
from typing import Any, Dict
from src.simulator.qccd_circuit import process_circuit, process_circuit_wise_arch
from src.simulator.color_code_processor import process_color_code_circuit, process_color_code_circuit_wise_arch, process_color_code_circuit_linear_arch, process_color_code_circuit_switch_arch
from datetime import datetime
from tqdm import tqdm
import json  

def get_logger(log_file: str) -> logging.Logger:
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    log_dir = os.path.dirname(log_file)
    if log_dir:
        os.makedirs(log_dir, exist_ok=True)
    handler = logging.FileHandler(log_file)
    formatter = logging.Formatter('%(processName)s - %(levelname)s - %(message)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    return logger

def save_results(data: Dict[str, Any], output_dir: str):
    """Save experiment results to a timestamped JSON file."""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    os.makedirs(output_dir, exist_ok=True)  
    output_path = os.path.join(output_dir, f"{timestamp}_experiment.json")

    with open(output_path, 'w') as outfile:
        json.dump(data, outfile, indent=4)

    print(f"Results saved to {output_path}")

def main(config_path: str):
    with open(config_path, 'r') as config_file:
        config = yaml.safe_load(config_file)

    hardware = config["hardware"]
    qec = config["qec"]
    simulation = config["simulation"]

    distances = qec["distances"]
    capacities = hardware["trap_capacity"]
    topology = hardware["topology"]
    gate_improvements = qec["gate_improvements"]
    num_shots = simulation["num_shots"]
    rounds = simulation["rounds"]
    num_cores = hardware.get("num_cores", os.cpu_count())
    placement_strategy = hardware.get("placement_strategy", "hill_climb")
    code_type = qec.get("code_type", "color_code")
    logger = get_logger(simulation["log_file"])
    data: Dict[str, Dict[str, Dict[int, Dict[int, Any]]]] = {
        "ElapsedTime": {}, "Operations": {}, "MeanConcurrency": {}, 
        "QubitOperations": {}, "LogicalErrorRates": {}, 
        "PhysicalZErrorRates": {}, "PhysicalXErrorRates": {}
    }

    logger.info("Starting parallel processing of circuits")
    with concurrent.futures.ProcessPoolExecutor(max_workers=num_cores) as executor:
        if code_type == "surface_code":
            futures = [
                executor.submit(process_circuit, d, c, gate_improvements, num_shots)
                for d in distances for c in capacities
            ]
        elif code_type == "color_code":
            if topology == "grid":
                futures = [
                    executor.submit(process_color_code_circuit, d, rounds, c, gate_improvements, num_shots, (6,6,6), placement_strategy=placement_strategy)
                    for d in distances for c in capacities
                ]
            elif topology == "linear":
                futures = [
                    executor.submit(process_color_code_circuit_linear_arch, d, rounds, c, gate_improvements, num_shots, (6,6,6))
                    for d in distances for c in capacities
                ]
            elif topology == "switch":
                futures = [
                    executor.submit(process_color_code_circuit_switch_arch, d, rounds, c, gate_improvements, num_shots, (6,6,6))
                    for d in distances for c in capacities
                ]
            else:
                raise ValueError(f"Unsupported topology for color_code: {topology!r}")
        else:
            raise ValueError(f"Unsupported code_type: {code_type!r}")

        pbar = tqdm(total=(len(distances)*len(capacities)*len(gate_improvements)))

        for future in concurrent.futures.as_completed(futures):
            try:
                result = future.result()
                d = result["Distance"]
                c = result["Capacity"]
                for metric in data.keys():
                    if metric in result:
                        for label, value in result[metric].items():
                            data[metric].setdefault(label, {}).setdefault(d, {})[c] = value
                """            
                for label in result["ElapsedTime"]:
                    data["ElapsedTime"][label][d][c] = result["ElapsedTime"][label]
                    data["Operations"][label][d][c] = result["Operations"][label]
                    data["MeanConcurrency"][label][d][c] = result["MeanConcurrency"][label]
                    data["QubitOperations"][label][d][c] = result["QubitOperations"][label]
                    data["LogicalErrorRates"][label][d][c] = result["LogicalErrorRates"][label]
                    data["PhysicalXErrorRates"][label][d][c] = result["PhysicalXErrorRates"][label]
                    data["PhysicalZErrorRates"][label][d][c] = result["PhysicalZErrorRates"][label]"""
                logger.info(f"Processed results for distance {d}, capacity {c}.")
                pbar.update(1)
            except Exception as e:
                logger.error("An error occurred during processing", exc_info=e)
        pbar.close()
    save_results(data, output_dir=simulation.get("output_dir", "data"))

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Run QCCD circuit simulations.")
    parser.add_argument("--config", type=str, required=True, help="Path to the configuration YAML file.")
    args = parser.parse_args()
    main(args.config)