"""Torch C++ contraction backend for cotengra trees.

This module compiles a restricted cotengra contraction tree into a compact
execution plan, then executes that plan in a C++ torch extension. It is an
experimental CPU path for reducing Python-level overhead between many
pairwise contractions.
"""

from __future__ import annotations

import importlib
import os
from functools import lru_cache
from pathlib import Path
from collections import defaultdict


_EXTENSION = None
_CONTRACTORS = {}
SMALL_GEMM_BATCH_FLOPS = 1_000_000


def _load_extension():
    global _EXTENSION
    if _EXTENSION is not None:
        return _EXTENSION

    from torch.utils.cpp_extension import load

    source = Path(__file__).resolve().parent / "csrc" / "torch_contractor.cpp"
    mklroot = os.environ.get("MKLROOT")
    extra_cflags = ["-O3"]
    extra_ldflags = []
    extra_include_paths = []
    if mklroot:
        mklroot_path = Path(mklroot)
        mkl_include = mklroot_path / "include"
        mkl_lib = mklroot_path / "lib"
        if (mkl_include / "mkl_cblas.h").exists() and (
            (mkl_lib / "libmkl_rt.so").exists()
            or (mkl_lib / "libmkl_rt.so.2").exists()
        ):
            extra_cflags.append("-DQIBOTN_USE_MKL")
            extra_include_paths.append(str(mkl_include))
            extra_ldflags.extend([f"-L{mkl_lib}", "-lmkl_rt"])

    _EXTENSION = load(
        name="qibotn_torch_contractor",
        sources=[str(source)],
        extra_cflags=extra_cflags,
        extra_ldflags=extra_ldflags,
        extra_include_paths=extra_include_paths,
        verbose=False,
    )
    return _EXTENSION


def _is_plain_permutation(expr):
    if expr is None:
        return None
    if isinstance(expr, tuple):
        return tuple(int(i) for i in expr)
    if not isinstance(expr, str):
        return None
    if "," in expr or "->" not in expr:
        return None
    source, target = expr.split("->", 1)
    if len(source) != len(target):
        return None
    if len(set(source)) != len(source) or set(source) != set(target):
        return None
    return tuple(source.index(ix) for ix in target)


def _maybe_tuple(values):
    return () if values is None else tuple(int(x) for x in values)


def _shape_from_inds(tree, node):
    return tuple(int(tree.size_dict[ix]) for ix in tree.get_inds(node))


def _matmul_signature(op):
    kind = op[3]
    if kind != 0:
        return None
    left_shape = op[5]
    right_shape = op[7]
    if len(left_shape) == 2 and len(right_shape) == 2:
        m, k, n = left_shape[-2], left_shape[-1], right_shape[-1]
        return ("mm", int(m), int(k), int(n), int(m * k * n))
    return None


def _normalize_node_ids(tree, contractions):
    leaf_to_id = {
        frozenset((i,)): i
        for i in range(tree.N)
    }
    next_id = len(leaf_to_id)
    node_to_id = dict(leaf_to_id)
    for parent, _left, _right, _tdot, _arg, _perm in contractions:
        if parent not in node_to_id:
            node_to_id[parent] = next_id
            next_id += 1

    return node_to_id, next_id


@lru_cache(maxsize=32)
def compile_torch_plan(tree):
    """Compile ``tree`` into C++ contractor plan fields.

    The supported subset is the same pairwise matmul lowering used by
    cotengra for torch CPU. Single-tensor diagonal/sum preprocessing is not
    supported yet because it appears only in less common trees; callers should
    fall back to cotengra for those cases.
    """

    contract_mod = importlib.import_module("cotengra.contract")
    contractions = contract_mod.extract_contractions(tree)
    node_to_id, ntemps = _normalize_node_ids(tree, contractions)
    plan = []

    for parent, left, right, tdot, arg, perm in contractions:
        if left is None or right is None:
            raise NotImplementedError(
                "C++ torch contractor does not support cotengra preprocessing."
            )

        left_shape = _shape_from_inds(tree, left)
        right_shape = _shape_from_inds(tree, right)
        if tdot:
            parsed = contract_mod._parse_tensordot_axes_to_matmul(
                arg,
                left_shape,
                right_shape,
            )
        else:
            parsed = contract_mod._parse_eq_to_batch_matmul(
                arg,
                left_shape,
                right_shape,
            )

        (
            eq_a,
            eq_b,
            new_shape_a,
            new_shape_b,
            new_shape_ab,
            perm_ab,
            pure_multiplication,
        ) = parsed

        left_perm = _is_plain_permutation(eq_a)
        right_perm = _is_plain_permutation(eq_b)
        if left_perm is None and eq_a is not None:
            raise NotImplementedError(f"Unsupported left preparation: {eq_a!r}")
        if right_perm is None and eq_b is not None:
            raise NotImplementedError(f"Unsupported right preparation: {eq_b!r}")

        plan.append(
            (
                node_to_id[parent],
                node_to_id[left],
                node_to_id[right],
                1 if pure_multiplication else 0,
                left_perm or (),
                _maybe_tuple(new_shape_a),
                right_perm or (),
                _maybe_tuple(new_shape_b),
                _maybe_tuple(new_shape_ab),
                _maybe_tuple(perm_ab),
            )
        )

        if perm is not None:
            raise NotImplementedError(
                "C++ torch contractor does not support cotengra tensordot perm."
            )

    root_id = node_to_id[tree.root]
    return tuple(plan), int(ntemps), int(root_id)


@lru_cache(maxsize=32)
def compile_batch_groups(tree, max_flops=SMALL_GEMM_BATCH_FLOPS):
    plan, _ntemps, _root_id = compile_torch_plan(tree)
    contractions = importlib.import_module("cotengra.contract").extract_contractions(tree)
    node_to_id, _ntemps = _normalize_node_ids(tree, contractions)
    depth = {frozenset((i,)): 0 for i in range(tree.N)}
    tensor_depth = {i: 0 for i in range(tree.N)}
    groups = defaultdict(list)

    for op_index, (contract_op, contraction) in enumerate(zip(plan, contractions)):
        parent, left, right, _tdot, _arg, _perm = contraction
        d = max(depth[left], depth[right]) + 1
        depth[parent] = d
        tensor_depth[contract_op[0]] = d
        sig = _matmul_signature(contract_op)
        if sig is None:
            continue
        kind, m, k, n, flops = sig
        if flops > max_flops:
            continue
        groups[(d, kind, m, k, n)].append(op_index)

    batch_groups = tuple(
        tuple(items)
        for _key, items in sorted(groups.items(), key=lambda item: (item[0], item[1][0]))
        if len(items) >= 2
    )
    return batch_groups


def batch_group_summary(tree, max_flops=SMALL_GEMM_BATCH_FLOPS):
    plan, _ntemps, _root_id = compile_torch_plan(tree)
    groups = compile_batch_groups(tree, max_flops=max_flops)
    covered = sum(len(group) for group in groups)
    calls_saved = sum(len(group) - 1 for group in groups)
    by_shape = []
    for group in groups:
        op = plan[group[0]]
        sig = _matmul_signature(op)
        by_shape.append((sig[1:4], len(group), group[:8]))
    return {
        "groups": len(groups),
        "covered_ops": covered,
        "calls_saved": calls_saved,
        "by_shape": by_shape,
    }


def contract_tree_cpp(tree, arrays):
    """Contract a cotengra tree using the experimental C++ torch contractor."""

    contractor = prepare_torch_cpp_contractor(tree)
    return contractor.contract(list(arrays))


def prepare_torch_cpp_contractor(tree):
    """Load the extension and compile ``tree`` without running contraction."""

    ext = _load_extension()
    key = id(tree)
    contractor = _CONTRACTORS.get(key)
    if contractor is None:
        plan, ntemps, root_id = compile_torch_plan(tree)
        contractor = ext.Contractor(list(plan), ntemps, root_id)
        _CONTRACTORS[key] = contractor
    return contractor
