"""Fill in a module description here"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['get_parser', 'prepare_args', 'show_CX_block', 'random_qubit_pairs', 'generate_random_gateList', 'generate_random',
           'save_gateList', 'main']

# %% ../nbs/00_core.ipynb 3
import os, hashlib
import numpy as np
from time import time
from pprint import pprint
from typing import Tuple, Dict, Any

def get_parser():
    "Create argument parser for random CX-block circuit generator"
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("-v","--verb",type=int, help="increase debug verbosity", default=1)
    parser.add_argument('-q','--numQubits', default=5, type=int, help='circuit width')
    parser.add_argument('-k','--numCX', default=4, type=int, help='num of CX gates')
    parser.add_argument('-i','--numCirc', default=2, type=int, help='num of circuits in the job')
    parser.add_argument("--expName",  default=None, help='(optional) experiment name')
    parser.add_argument("--basePath", default=None, help="head path for experiments, or 'env'")
    parser.add_argument("--outPath",  default='out/', help="(optional) redirect outputs")
    return parser

# %% ../nbs/00_core.ipynb 4
def prepare_args(args):
    "Process arguments and ensure output directories exist"
    import os
    if args.basePath == 'env':
        args.basePath = os.environ['Cudaq_dataVault']
    if args.basePath and args.outPath:
        args.outPath = os.path.join(args.basePath, 'circ')
    os.makedirs(args.outPath, exist_ok=True)
    assert args.numQubits >= 2
    return args

# %% ../nbs/00_core.ipynb 5
def show_CX_block():
    "Display a CX-block circuit with parameters θ and φ"
    from qiskit import QuantumCircuit
    from qiskit.circuit import Parameter
    theta = Parameter('θ')
    phi = Parameter('φ')
    qc = QuantumCircuit(2)
    qc.barrier()
    qc.ry(theta, 0)
    qc.rz(phi, 1)
    qc.cx(0, 1)
    qc.barrier()
    print(qc.draw())

# %% ../nbs/00_core.ipynb 6
def random_qubit_pairs(nq:int, k:int):
    "Generate k random ordered qubit pairs from nq qubits"
    all_pairs = np.array([(i, j) for i in range(nq) for j in range(nq) if i != j])
    selected_indices = np.random.choice(len(all_pairs), k, replace=True)
    return all_pairs[selected_indices]

# %% ../nbs/00_core.ipynb 7
def generate_random_gateList(numQubits:int, numCX:int, numCirc:int):
    "Generate random CX-block circuits and return data, metadata"
    nGate = 3 * numCX
    m = {'h': 1, 'ry': 2, 'rz': 3, 'cx': 4, 'measure': 5}
    circ_type = np.zeros((numCirc, 2), dtype=np.int32)
    gate_type = np.zeros((numCirc, nGate, 3), dtype=np.int32)
    gate_param = np.random.uniform(0, np.pi, size=(numCirc, nGate)).astype(np.float32)
    t_ry = np.full((numCX, 1), m['ry'])
    t_rz = np.full((numCX, 1), m['rz'])
    t_cx = np.full((numCX, 1), m['cx'])

    for j in range(numCirc):
        qpairs = random_qubit_pairs(numQubits, numCX)
        rpairs = qpairs[:, ::-1]
        circ_type[j] = [numQubits, nGate]
        gate_type[j, 0::3] = np.concatenate((t_ry, qpairs), axis=1)
        gate_type[j, 1::3] = np.concatenate((t_rz, rpairs), axis=1)
        gate_type[j, 2::3] = np.concatenate((t_cx, qpairs), axis=1)
        gate_param[j, 2::3] = 0

    outD = {'circ_type': circ_type, 'gate_type': gate_type, 'gate_param': gate_param}
    md = {'gate_map': m, 'num_cx': numCX, 'num_qubit': numQubits,
          'num_gate': nGate, 'num_circ': numCirc}
    return outD, md

# %% ../nbs/00_core.ipynb 8
import os, hashlib
from toolbox.Util_H5io4 import write4_data_hdf5
from pprint import pprint

def generate_random(
    numCX:int = 4,
    numCirc:int = 1,
    numQubits:int = 5,
    expName:str = None,
    basePath:str = None,
    outPath:str = None
):
    """
    High-level wrapper: generate random CX-block circuits and save to HDF5.

    Args:
        numCX: Number of CX gates per circuit (default=4)
        numCirc: Number of circuits to generate (default=1)
        numQubits: Number of qubits in each circuit (default=5)
        expName: Optional experiment name; if None, a random hash is used
        basePath: Base directory for storing results; if None, defaults to `$Cudaq_dataVault`
        outPath: Optional override for output directory; if None, defaults to `<basePath>/circ`

    Returns:
        Path to the generated `.gate_list.h5` file
    """
    from qgear.core import generate_random_gateList 

    # Generate circuit data
    outD, md = generate_random_gateList(numQubits, numCX, numCirc)
    md['hash'] = hashlib.md5(os.urandom(32)).hexdigest()[:6]
    md['short_name'] = expName if expName else f"rcirc_{md['hash']}"

    # Resolve base path
    if basePath is None or basePath == 'env':
        basePath = os.environ.get('Cudaq_dataVault', os.getcwd())  # fallback to current dir
    # Resolve output path
    if outPath is None:
        outPath = os.path.join(basePath, 'circ')
    os.makedirs(outPath, exist_ok=True)

    # Save to HDF5
    outF = os.path.join(outPath, f"{md['short_name']}.gate_list.h5")
    write4_data_hdf5(outD, outF, md)
    pprint(md)
    return outF

# %% ../nbs/00_core.ipynb 9
def save_gateList(outD:Dict[str,Any], md:Dict[str,Any], outPath:str, expName:str=None):
    "Save generated gate list to HDF5 file"
    from Util_H5io4 import write4_data_hdf5
    md['hash'] = hashlib.md5(os.urandom(32)).hexdigest()[:6]
    md['short_name'] = expName if expName else f"rcirc_{md['hash']}"
    outF = os.path.join(outPath, f"{md['short_name']}.gate_list.h5")
    write4_data_hdf5(outD, outF, md)
    pprint(md)
    return outF

# %% ../nbs/00_core.ipynb 10
def main():
    parser = get_parser()
    args = parser.parse_args()
    args = prepare_args(args)
    show_CX_block()
    T0 = time()
    outD, md = generate_random_gateList(args.numQubits, args.numCX, args.numCirc)
    print(f"M:  gen_rand  elaT= {time()-T0:.1f} sec")
    save_gateList(outD, md, args.outPath, args.expName)
