#!/usr/bin/env python3
"""Inspect Core ML operation placement and estimated accelerator cost share."""

from __future__ import annotations

import argparse
import platform
import sys
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any, Sequence

from purpose_data import DataError, write_json


def _device_category(device: Any) -> str:
    name = type(device).__name__.lower()
    description = str(device).lower()
    combined = f"{name} {description}"
    if "neural" in combined:
        return "neuralEngine"
    if "gpu" in combined:
        return "gpu"
    if "cpu" in combined:
        return "cpu"
    return "unknown"


def _compute_unit(coremltools: Any, requested: str) -> Any:
    values = {
        "all": coremltools.ComputeUnit.ALL,
        "cpu-only": coremltools.ComputeUnit.CPU_ONLY,
        "cpu-and-gpu": coremltools.ComputeUnit.CPU_AND_GPU,
        "cpu-and-ne": coremltools.ComputeUnit.CPU_AND_NE,
    }
    return values[requested]


def inspect(args: argparse.Namespace) -> dict[str, Any]:
    if not args.model.exists():
        raise DataError(f"{args.model}: Core ML model is missing")
    try:
        import coremltools as ct
    except ImportError as exc:
        raise DataError(
            "Core ML inspection requires requirements-coreml.txt on macOS"
        ) from exc

    compiled = ct.models.utils.compile_model(str(args.model))
    compute_plan = ct.models.compute_plan.MLComputePlan.load_from_path(
        path=str(compiled),
        compute_units=_compute_unit(ct, args.compute_units),
    )
    program = compute_plan.model_structure.program
    if program is None or "main" not in program.functions:
        raise DataError("Core ML package is not an ML Program with a main function")
    operations = list(program.functions["main"].block.operations)
    if not operations:
        raise DataError("Core ML compute plan contains no operations")

    preferred_counts: Counter[str] = Counter()
    preferred_costs: dict[str, float] = defaultdict(float)
    supported_counts: Counter[str] = Counter()
    operation_reports = []
    operations_with_usage = 0
    operations_with_cost = 0
    total_cost = 0.0
    for operation in operations:
        usage = compute_plan.get_compute_device_usage_for_mlprogram_operation(
            operation
        )
        cost = compute_plan.get_estimated_cost_for_mlprogram_operation(operation)
        preferred = "unknown"
        supported: list[str] = []
        if usage is not None:
            operations_with_usage += 1
            preferred = _device_category(usage.preferred_compute_device)
            preferred_counts[preferred] += 1
            supported = sorted(
                {_device_category(device) for device in usage.supported_compute_devices}
            )
            supported_counts.update(supported)
        weight = None
        if cost is not None:
            operations_with_cost += 1
            weight = float(cost.weight)
            total_cost += weight
            preferred_costs[preferred] += weight
        operation_reports.append(
            {
                "operatorName": str(operation.operator_name),
                "preferredDevice": preferred,
                "supportedDevices": supported,
                "estimatedCostWeight": weight,
            }
        )

    ane_operations = preferred_counts["neuralEngine"]
    ane_cost = preferred_costs["neuralEngine"]
    report = {
        "schemaVersion": 1,
        "model": str(args.model),
        "coremltoolsVersion": ct.__version__,
        "machine": platform.machine(),
        "macOS": platform.mac_ver()[0],
        "computeUnits": args.compute_units,
        "operations": len(operations),
        "operationsWithDeviceUsage": operations_with_usage,
        "operationsWithEstimatedCost": operations_with_cost,
        "preferredOperationCounts": dict(sorted(preferred_counts.items())),
        "supportedOperationCounts": dict(sorted(supported_counts.items())),
        "preferredEstimatedCosts": dict(sorted(preferred_costs.items())),
        "neuralEngineOperationShare": (
            ane_operations / operations_with_usage if operations_with_usage else 0.0
        ),
        "neuralEngineEstimatedCostShare": (
            ane_cost / total_cost if total_cost else 0.0
        ),
        "operationDetails": operation_reports,
    }
    write_json(args.report, report)
    print(
        "Core ML placement: "
        f"ANE operations={report['neuralEngineOperationShare']:.2%} "
        f"ANE estimated cost={report['neuralEngineEstimatedCostShare']:.2%}"
    )
    return report


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--report", type=Path, required=True)
    parser.add_argument(
        "--compute-units",
        choices=("all", "cpu-only", "cpu-and-gpu", "cpu-and-ne"),
        default="cpu-and-ne",
    )
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        inspect(args)
    except (DataError, OSError, RuntimeError, ValueError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())