import json
import re
import shutil
import struct
from pathlib import Path

import numpy as np


_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 read_raptor_instruction_count(path):
    with open(path, "rb") as f:
        magic, version, count = struct.unpack("<4sII", f.read(12))
    if magic != b"PIMB" or version != 1:
        raise ValueError(f"Invalid Raptor binary instruction file: {path}")
    return count


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")

    binary_sources = sorted(pim_dir.glob("core_*.pim"), key=lambda path: int(path.stem.split("_")[1]))
    if binary_sources:
        record_dtype = np.dtype([
            ("opcode", "u1"), ("rd", "u1"), ("r1", "u1"), ("flags", "u1"),
            ("r2_or_imm", "<i4"), ("generic1", "<i4"),
            ("generic2", "<i4"), ("generic3", "<i4"),
        ])
        for source in binary_sources:
            destination = output_dir / source.name
            shutil.copyfile(source, destination)
            count = read_raptor_instruction_count(destination)
            if destination.stat().st_size != 12 + count * record_dtype.itemsize:
                raise ValueError(f"Invalid Raptor binary instruction file: {source}")
            records = np.memmap(destination, dtype=record_dtype, mode="r+", offset=12, shape=(count,))
            scalar_sldi = np.zeros(count, dtype=bool)
            vmv = records["opcode"] == 22
            sldi = records["opcode"] == 1
            for register in np.unique(records["r2_or_imm"][vmv]):
                definitions = np.flatnonzero(sldi & (records["rd"] == register))
                consumers = np.flatnonzero(vmv & (records["r2_or_imm"] == register))
                reaching = np.searchsorted(definitions, consumers) - 1
                scalar_sldi[definitions[reaching[reaching >= 0]]] = True
            for start in range(0, count, 1_000_000):
                chunk = records[start:start + 1_000_000]
                setbw = chunk["opcode"] == 8
                chunk["generic1"][setbw] = 8
                chunk["generic2"][setbw] = 8
                address_sldi = (chunk["opcode"] == 1) & ~scalar_sldi[start:start + len(chunk)]
                if np.any(chunk["r2_or_imm"][address_sldi] % 4):
                    raise ValueError(f"Raptor address is not aligned to its fp32 element width: {source}")
                chunk["r2_or_imm"][address_sldi] //= 4
                transfer = (chunk["opcode"] >= 25) & (chunk["opcode"] <= 30)
                if np.any(chunk["generic2"][transfer] % 4) or np.any(chunk["generic3"][transfer] % 4):
                    raise ValueError(f"Raptor transfer field is not aligned to its fp32 element width: {source}")
                chunk["generic2"][transfer] //= 4
                chunk["generic3"][transfer] //= 4
            records.flush()
            del records
        return output_dir

    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)
        last_sldi = {}
        scalar_sldi = set()
        for index, instruction in enumerate(instructions):
            if instruction["op"] == "sldi":
                last_sldi[instruction["rd"]] = index
            elif instruction["op"] == "vmv" and instruction["rs2"] in last_sldi:
                scalar_sldi.add(last_sldi[instruction["rs2"]])
        for index, instruction in enumerate(instructions):
            op = instruction["op"]
            if op == "setbw":
                instruction["ibiw"] = 8
                instruction["obiw"] = 8
            elif op == "sldi" and index not in scalar_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