# Copyright 2025-2026 XMOS LIMITED.
# This Software is subject to the terms of the XMOS Public Licence: Version 1.

from pathlib import Path
import itertools
import subprocess
import re
import json

def max_mips(lines):
    mips_values = []
    for line in lines:
        # Look for patterns like "12.1539 MIPS"
        match = re.search(r"([\d.]+)\s*MIPS", line)
        if match:
            mips_values.append(float(match.group(1)))

    return max(mips_values) if mips_values else None

def write_rst_table(results: dict, outfile: Path):
    """
    Write results dict to an RST list-table.
    cfg key format: '{mics}mic_{pdmrx}_{fs}fs'
    """
    rows = []
    for cfg, mips in sorted(results.items()):
        # Parse cfg
        # e.g. 2mic_thread_32000fs
        try:
            mic_part, pdmrx_part, fs_part = cfg.split('_')
            mic_count = mic_part.replace('mic', '')  # '2mic' -> '2'
            pdmrx_mode = pdmrx_part
            fs_val = fs_part.replace('fs', '')
        except ValueError:
            continue
        mips_str = f"{mips:.3f}" if isinstance(mips, (int, float)) else "N/A"
        rows.append((mic_count, pdmrx_mode.upper(), fs_val, mips_str))

    lines = []
    lines.append(".. _mic_array_mips:\n")
    lines.append(".. list-table:: Estimated MIPS (per configuration)")
    lines.append("   :header-rows: 1")
    lines.append("   :widths: 6 6 8 8")
    lines.append("")
    lines.append("   * - mic count")
    lines.append("     - PDM RX")
    lines.append("     - output samp freq")
    lines.append("     - MIPS")
    for mic_count, pdmrx_mode, fs_val, mips_str in rows:
        lines.append(f"   * - {mic_count}")
        lines.append(f"     - {pdmrx_mode}")
        lines.append(f"     - {fs_val}")
        lines.append(f"     - {mips_str}")
    outfile.write_text("\n".join(lines))

def test_measure_mips(pytestconfig):
    """
    Run MIPS-usage measurements for all microphone array configurations and
    compare them with stored reference data.

    Behaviour:
      • When --update is not passed:
          The measured MIPS values are checked against the reference JSON file
          (mic_array_mips.json). The test fails if any configuration deviates
          from the reference by more than ±0.05 MIPS. This ensures the mic array MIPS
          remain stable across changes.

      • When --update is passed:
          The reference JSON file is overwritten with the newly measured
          values. mic_array_mips_table.rst is also regenerated.

    Files:
      mic_array_mips.json         - stored reference MIPS values
      mic_array_mips_table.rst    - autogenerated RST table of results
    """
    update = pytestconfig.getoption("--update")
    cwd = Path(__file__).parent
    mics = [1, 2]
    pdmrx = ["isr", "thread"]
    fs = [16000, 32000, 48000]
    results = {}
    for chans, pdmrx_type, samp_freq in itertools.product(mics, pdmrx, fs):
        cfg = f"{chans}mic_{pdmrx_type}_{samp_freq}fs"
        xe_path = f'{cwd}/app_mips/bin/{cfg}/test_mips_{cfg}.xe'
        assert Path(xe_path).exists(), f"Cannot find {xe_path}"
        ret = subprocess.run(["xrun", "--xscope", "--id", "0", xe_path], capture_output=True, text=True, check=True, timeout=15)
        results[cfg] = max_mips(ret.stdout.splitlines())

    # Compare against mic_array_mips.json that's already there to ensure MIPS
    # number are in the same ballpark, before overwriting mic_array_mips.json
    outfile = cwd / "mic_array_mips.json"
    with outfile.open("r") as f:
        ref_data = json.load(f)
        for cfg in ref_data:
            ref_mips = ref_data[cfg]
            assert cfg in results, f"cfg {cfg} not found in results.\nresults = {results}"
            test_mips = results[cfg]
            if not update:
                threshold = 0.05
                assert abs(test_mips - ref_mips) < threshold, (f"For cfg {cfg}, test_mips {test_mips} differ "
                                                               f"from ref_mips {ref_mips} by more than the allowed threshold of {threshold}.\n"
                                                               f"If this is expected, run test with pytest test_measure_mips --update "
                                                               f"to update the reference json file and the table rst.\n")

    if update:
        with outfile.open("w") as f:
            json.dump(results, f, indent=2)

        # RST table output
        rst_out = cwd / "mic_array_mips_table.rst"
        write_rst_table(results, rst_out)


