import json
import re
import shutil
from pathlib import Path


_METRIC_PATTERNS = {
    "output_count": r"output count:\s+([0-9]+)\s+samples",
    "throughput": r"throughput:\s+([0-9.eE+-]+)\s+samples/s",
    "average_latency_ms": r"average latency:\s+([0-9.eE+-]+)\s+ms",
    "latency_ms": r"latency:\s+([0-9.eE+-]+)\s+ms",
    "average_power_mw": r"average power:\s+([0-9.eE+-]+)\s+mW",
    "average_energy_pj": r"average energy:\s+([0-9.eE+-]+)\s+pJ(?:/it)?",
}


def parse_pimsim_nn_metrics(output):
    metrics = {"raw_output": output}
    for name, pattern in _METRIC_PATTERNS.items():
        match = re.search(pattern, output)
        if match:
            value = match.group(1)
            metrics[name] = int(value) if name == "output_count" else float(value)
    return metrics


def export_raptor_latency_artifact(pim_dir, output_dir):
    pim_dir = Path(pim_dir)
    output_dir = Path(output_dir)
    if output_dir.exists():
        shutil.rmtree(output_dir)
    output_dir.mkdir(parents=True)

    def int8_bytes(value, field):
        if value % 4:
            raise ValueError(f"Raptor {field}={value} is not aligned to its fp32 element width")
        return value // 4

    with open(pim_dir / "config.json", encoding="utf-8") as f:
        config = json.load(f)
    for field in ("inputs_addresses", "outputs_addresses"):
        if field in config:
            config[field] = [int8_bytes(value, field) for value in config[field]]
    with open(output_dir / "config.json", "w", encoding="utf-8") as f:
        json.dump(config, f, separators=(",", ":"))
        f.write("\n")

    byte_size_fields = {
        "ld": "size",
        "st": "size",
        "lldi": "len",
        "lmv": "len",
        "send": "size",
        "recv": "size",
    }
    for source in sorted(pim_dir.glob("core_*.json"), key=lambda path: int(path.stem.split("_")[1])):
        with open(source, encoding="utf-8") as f:
            instructions = json.load(f)
        for instruction in instructions:
            op = instruction["op"]
            if op == "setbw":
                instruction["ibiw"] = 8
                instruction["obiw"] = 8
            elif op == "sldi":
                instruction["imm"] = int8_bytes(instruction["imm"], "address")
            if field := byte_size_fields.get(op):
                instruction[field] = int8_bytes(instruction[field], f"{op} {field}")
                if offset := instruction.get("offset"):
                    offset["offset_value"] = int8_bytes(offset["offset_value"], f"{op} offset")
        with open(output_dir / source.name, "w", encoding="utf-8") as f:
            json.dump(instructions, f, separators=(",", ":"))
            f.write("\n")
    return output_dir