# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_plotter.ipynb.

# %% auto 0
__all__ = ['extract_date_from_path', 'extract_qft_from_filename', 'readOne', 'readOneQFT', 'find_yaml_files', 'sort_end_lists',
           'MetricsPlotter', 'metrics_plot', 'postproc_qcrank', 'restore_canned_image', 'residual_analysis',
           'process_qcrank_experiment']

# %% ../nbs/02_plotter.ipynb 1
import os, numpy as np
from .toolbox.PlotterBackbone import PlotterBackbone
from .toolbox.Util_IOfunc import read_yaml
from pprint import pprint
from time import time
from .toolbox.Util_H5io4 import write4_data_hdf5, read4_data_hdf5
from .toolbox.Util_ibm import marginalize_qcrank_EV
from .toolbox.Util_Qiskit import unpack_numpy_to_counts
from .toolbox.PlotterQCrank import Plotter

def extract_date_from_path(file_path):
    "Extract date from a path containing 'dataCudaQ_<DATE>'"
    for component in file_path.split('/'):
        if component.startswith('dataCudaQ_'):
            return component[len('dataCudaQ_'):]
    return None

def extract_qft_from_filename(filename):
    "Check if filename contains 'qft1' and return it if found"
    for component in filename.split('_'):
        if component == "qft1":
            return component
    return None

# %% ../nbs/02_plotter.ipynb 2
def readOne(inpF, dataD, verb=1):
    "Read one YAML metrics file (non-QFT)"
    assert os.path.exists(inpF)
    date = extract_date_from_path(inpF)
    xMD = read_yaml(inpF, verb)
    nq = float(xMD['num_qubit'])
    runt = float(xMD['elapsed_time']) / float(xMD['num_circ'])
    cores = xMD.get('cores', 32)
    tasks_per_node = xMD.get('tasks_per_node', 4)
    tag1 = 'cpu' if 'cpu_info' in xMD else 'gpu'
    tag2 = xMD['target']
    if tag1 not in dataD:
        dataD[tag1] = {}
    num_cx_formatted = "10k" if xMD["num_cx"] == 10000 else f'{xMD["num_cx"]}'
    if tag1 == 'cpu':
        tag3 = f'{num_cx_formatted}CX_c{cores}_tp{tasks_per_node}'
    else:
        g_tag = tag2.split('-')[1] if '-' in tag2 else 'gpu'
        tag3 = f'{g_tag}.{num_cx_formatted}CX'
    if tag2 not in dataD[tag1]:
        dataD[tag1][tag2] = {}
    if tag3 not in dataD[tag1][tag2]:
        dataD[tag1][tag2][tag3] = {'nq': [], 'runt': [], 'cores': [], 'tasks_per_node': [], 'date': []}
    head = dataD[tag1][tag2][tag3]
    head['nq'].append(nq)
    head['runt'].append(runt)
    head['cores'].append(cores)
    head['tasks_per_node'].append(tasks_per_node)
    head['date'].append(date)

# %% ../nbs/02_plotter.ipynb 3
def readOneQFT(inpF, dataD, qft, verb=1):
    "Read one YAML metrics file for QFT runs"
    assert os.path.exists(inpF)
    date = extract_date_from_path(inpF)
    xMD = read_yaml(inpF, verb)
    nq = float(xMD['num_qubit'])
    runt = float(xMD['elapsed_time'])
    tag1 = qft
    tag2 = xMD['target']
    if tag1 not in dataD:
        dataD[tag1] = {}
    num_shots_formatted = "10k" if xMD["num_shots"] == 10000 else xMD["num_shots"]
    options = inpF.split('/')[-1].split('_')[-2]
    g_tag = tag2.split('-')[1] if '-' in tag2 else 'gpu'
    tag3 = f'{qft}.{g_tag}.{options}.{num_shots_formatted}S'
    if tag2 not in dataD[tag1]:
        dataD[tag1][tag2] = {}
    if tag3 not in dataD[tag1][tag2]:
        dataD[tag1][tag2][tag3] = {'nq': [], 'runt': [], 'shots': [], 'date': []}
    head = dataD[tag1][tag2][tag3]
    head['nq'].append(nq)
    head['runt'].append(runt)
    head['shots'].append(xMD['num_shots'])
    head['date'].append(date)

# %% ../nbs/02_plotter.ipynb 4
def find_yaml_files(directory_path, vetoL=None):
    "Find all .yaml files in a directory, excluding files with veto strings"
    if vetoL is None:
        vetoL = []
    yaml_files = []
    for root, dirs, files in os.walk(directory_path):
        for file in files:
            if file.endswith('.yaml') and not any(veto in file for veto in vetoL):
                yaml_files.append(os.path.join(root, file))
    return yaml_files

# %% ../nbs/02_plotter.ipynb 5
def sort_end_lists(d, sort_key='nq', val_key='runt', parent_key=''):
    "Recursively sort lists in nested dicts by a sort_key"
    if sort_key in d:
        xV, yV = d[sort_key], d[val_key]
        xU, yU = map(list, zip(*sorted(zip(xV, yV), key=lambda x: x[0])))
        d[sort_key] = np.array(xU)
        d[val_key] = np.array(yU)
        return
    for k, v in d.items():
        if isinstance(v, dict):
            sort_end_lists(v, sort_key, val_key, f"{parent_key}.{k}" if parent_key else k)

# %% ../nbs/02_plotter.ipynb 6
class MetricsPlotter(PlotterBackbone):
    "Plotter for concatenated metrics"
    def __init__(self, prjName='metrics', shift=True, outPath='out', noXterm=True, verb=1):
        """
        Args:
            prjName: Project name for plots
            shift: Whether to shift CPU points to avoid overlap
            outPath: Output directory for plots
            noXterm: Disable X-term (for headless plotting)
            verb: Verbosity level
        """
        from types import SimpleNamespace
        args = SimpleNamespace(
            prjName=prjName,
            shift=shift,
            outPath=outPath,
            noXterm=noXterm,
            verb=verb
        )
        super().__init__(args)

    def compute_time(self, bigD, tag1, figId=1, shift=False):
        import numpy as np
        nrow, ncol = 1, 1
        figId = self.smart_append(figId)
        fig = self.plt.figure(figId, facecolor='white', figsize=(5.5, 7))
        ax = self.plt.subplot(nrow, ncol, 1)

        if 'gpu' in tag1:
            dataD = bigD[tag1.split('-')[1]]
        elif 'qft' in tag1:
            dataD = bigD[tag1]
        else:
            dataD = bigD[tag1]

        for tag2 in dataD:
            for tag3 in dataD[tag2]:
                if '20000CX' in tag3:
                    continue
                dataE = dataD[tag2][tag3]
                nqV = dataE['nq']
                runtV = dataE['runt'] / 60.0
                dLab = f'{tag3}'
                dCol = 'k'
                marker_style = 'o'

                if shift and tag1 == 'cpu':
                    shift_x = np.random.uniform(-0.1, 0.1, size=len(nqV))
                    shift_y = np.random.uniform(-0.1, 0.1, size=len(runtV))
                    ax.plot(nqV + shift_x, runtV + shift_y, marker=marker_style,
                            linestyle='-', color=dCol, label=dLab, markersize=9)
                else:
                    ax.plot(nqV, runtV, marker=marker_style, linestyle='-',
                            color=dCol, label=dLab, markersize=9)

        ax.set(xlabel='num qubits', ylabel='compute end-state (minutes)')
        ax.set_title(f'Compute state-vector tag1={tag1}', pad=50)
        ax.set_yscale('log')
        ax.set_ylim(1e-3, 1e+0)
        ax.set_xlim(15.5, 33.5)
        ax.grid()
        ax.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3,
                  ncol=3, mode="expand", borderaxespad=0., fontsize=8.5)

# %% ../nbs/02_plotter.ipynb 7
def metrics_plot(
    measPath: str = None,
    pathL=None,
    vetoL=None,
    showPlots='b',
    shift=True,
    outPath='out',
    noXterm=True,
    verb=1
):
    """
    Concatenate and plot metrics from multiple job YAML files.

    Args:
        corePath: Base path to data folders. Defaults to the current working directory.
        pathL: List of date folder suffixes (e.g., ['Nov15']). If None, auto-detects all `dataCudaQ_*` folders.
        vetoL: List of substrings to exclude from filenames.
        showPlots: String of plots to show ('a'=cpu, 'b'=par-gpu, 'c'=adj-gpu, 'd'=qft).
        shift: Whether to randomly shift CPU points to avoid overlap.
        outPath: Output directory for plots.
        noXterm: Disable X-term (for headless plotting).
        verb: Verbosity level.
    """
    import os
    if measPath is None:
        measPath = os.path.join(os.getcwd(), "meas")
    if vetoL is None:
        vetoL = []

    if not os.path.isdir(measPath):
        raise FileNotFoundError(f"Metrics directory not found: {measPath}")
    # Ensure output directory exists
    os.makedirs(outPath, exist_ok=True)
    # --- Scan meas directory for YAML files ---
    fileL = find_yaml_files(measPath, vetoL)
    if not fileL:
        raise FileNotFoundError(f"No YAML files found in {measPath}")

    dataAll, dataQFT = {}, {}
    for fileN in fileL:
        qft = extract_qft_from_filename(fileN)
        if not qft:
            readOne(fileN, dataAll)
        else:
            readOneQFT(fileN, dataQFT, qft)

    sort_end_lists(dataAll)
    sort_end_lists(dataQFT)

    plot = MetricsPlotter(prjName='metrics', shift=shift, outPath=outPath, noXterm=noXterm, verb=verb)
    if 'a' in showPlots:
        plot.compute_time(dataAll, 'cpu', figId=1, shift=shift)
    if 'b' in showPlots:
        plot.compute_time(dataAll, 'par-gpu', figId=2, shift=shift)
    if 'c' in showPlots:
        plot.compute_time(dataAll, 'adj-gpu', figId=3, shift=shift)
    if 'd' in showPlots:
        plot.compute_time(dataQFT, 'qft', figId=4, shift=shift)
    plot.display_all(png=1)

# %% ../nbs/02_plotter.ipynb 8
"""
Post-processing and plotting for QCrank experiments.
"""
# ----------------------------
def postproc_qcrank(expD, md, verb=1):
    """
    Post-process QCrank experiment results to reconstruct user data.

    Args:
        expD: Experiment data dictionary from HDF5.
        md: Metadata dictionary from HDF5.
        verb: Verbosity level.
    """
    pmd = md['payload']
    nq_addr = pmd['nq_addr']
    nq_data = pmd['nq_fdata']
    seq_len = pmd['seq_len']
    nImg = pmd['num_sample']

    countsL = unpack_numpy_to_counts(md, expD)

    rec_udata = np.zeros((nImg, nq_data, seq_len))  # reconstructed user data
    addrBitsL = [nq_data + i for i in range(nq_addr)]

    if verb:
        print('rec_udata:', rec_udata.shape, 'addrBitsL:', addrBitsL)

    T0 = time()
    for ic in range(nImg):
        counts = countsL[ic]
        for ibit in range(nq_data):
            T1 = time()
            rec_udata[ic, ibit] = marginalize_qcrank_EV(addrBitsL, counts, dataBit=ibit)
            if ic < 5 and verb:
                print(f"ic={ic} marginal ibit={ibit} done, elaT={(T1-T0)/60:.1f} min")

    expD['rec_udata'] = rec_udata

# %% ../nbs/02_plotter.ipynb 9
# ----------------------------
def restore_canned_image(expD, md):
    """
    Restore reconstructed image from QCrank output.

    Args:
        expD: Experiment data dictionary.
        md: Metadata dictionary.
    """
    cad = md['canned']
    n_img = md['payload']['num_sample']
    pixX, pixY = cad['image_shape_xy']

    assert n_img == 1
    recA = expD['rec_udata'][0]
    expD['rec_norm_image'] = recA.reshape(pixY, pixX)

# %% ../nbs/02_plotter.ipynb 10
# ----------------------------
def residual_analysis(expD, md):
    """
    Compute residual statistics between reconstructed and true data.

    Args:
        expD: Experiment data dictionary.
        md: Metadata dictionary.
    """
    rdata = expD['rec_udata'].flatten()
    tdata = expD['true_out_udata'].flatten()
    res_data = rdata - tdata
    mean = np.mean(res_data)
    std = np.std(res_data)
    N = res_data.shape[0]
    se_s = std / np.sqrt(2 * (N - 1))

    if 'postproc' not in md:
        md['postproc'] = {}
    md['postproc'].update({
        'res_mean': float(mean),
        'res_std': float(std),
        'res_SE_s': float(se_s)
    })

# %% ../nbs/02_plotter.ipynb 11
import os
import numpy as np
from pprint import pprint

def process_qcrank_experiment(
    exp_name,
    inp_path=None,
    out_path="out",
    show_plots="a",
    no_xterm=True,
    verb=1,
    save_plots=True
):
    """
    Load, post-process, and optionally plot QCrank experiment results.

    Args:
        exp_name: Experiment name (without .h5 extension).
        inp_path: Path to raw outputs from experiment. Defaults to current working directory.
        out_path: Path to save post-processed results.
        show_plots: String of plots to show ('a'=accuracy, 'b'=image, 'c'=dynamic range).
        no_xterm: Disable X-term for headless plotting.
        verb: Verbosity level.
        save_plots: If True, save plots to file.

    Returns:
        (expD, expMD) - processed data and metadata.
    """
    # Default to current working directory if inp_path is not provided
    if inp_path is None:
        inp_path = os.path.join(os.getcwd(), "out")

    np.set_printoptions(precision=3)

    inpF = exp_name + '.h5'
    expD, expMD = read4_data_hdf5(os.path.join(inp_path, inpF))
    if verb >= 2:
        print('M:expMD:')
        pprint(expMD)

    postproc_qcrank(expD, expMD, verb=verb)
    restore_canned_image(expD, expMD)

    # Save post-processed output
    os.makedirs(out_path, exist_ok=True)
    outF = os.path.join(out_path, expMD['short_name'] + '.post.h5')
    write4_data_hdf5(expD, outF, expMD)

    # Plotting
    expMD['plot'] = {'resid_max_range': 0.4}
    args = type("Args", (), {})()
    args.prjName = expMD['short_name']
    args.noXterm = no_xterm
    args.outPath = out_path
    args.verb = verb
    plot = Plotter(args)

    if 'a' in show_plots:
        plot.qcrank_accuracy(expD, expMD, figId=1)
    if 'b' in show_plots:
        plot.canned_image(expD, expMD, figId=2)
    if 'c' in show_plots:
        plot.dynamic_range(expD, expMD, figId=3)

    plot.display_all(png=int(save_plots))
    return expD, expMD
