"""Parallel path search and contraction utilities for tensor networks."""
import os
import pickle
import signal
import time
from math import log2, log10
import numpy as np
from dataclasses import dataclass
from concurrent.futures import ProcessPoolExecutor, TimeoutError, as_completed

try:
    from mpi4py import MPI
    _HAVE_MPI = True
except ImportError:
    _HAVE_MPI = False
    MPI = None


SEARCH_METHODS = ("greedy", "kahypar", "kahypar-agglom", "spinglass")
_COTENGRA_DASK_PATCHED = False
_COTENGRA_DASK_SUBMIT_PATCHED = False
_DASK_TRIAL_DEBUG = False


def _optimizer_search_stats(opt):
    scores = list(getattr(opt, "scores", ()))
    finite_scores = [score for score in scores if np.isfinite(score)]
    times = list(getattr(opt, "times", ()))
    best = getattr(opt, "best", {}) or {}
    return {
        "completed_trials": len(scores),
        "finite_trials": len(finite_scores),
        "failed_trials": len(scores) - len(finite_scores),
        "requested_trials": int(getattr(opt, "max_repeats", 0) or 0),
        "trial_seconds_sum": float(sum(times)),
        "best_score": float(best.get("score", float("inf"))),
        "best_flops": float(best.get("flops", float("inf"))),
        "best_write": float(best.get("write", float("inf"))),
        "best_size": float(best.get("size", float("inf"))),
    }


def _attach_search_stats(tree, opt):
    try:
        tree.qibotn_search_stats = _optimizer_search_stats(opt)
    except Exception:
        pass
    return tree


def _dask_worker_slots(client):
    info = client.scheduler_info(n_workers=-1)
    workers = info.get("workers", {})
    return workers, sum(int(w.get("nthreads", 1) or 1) for w in workers.values())


def _print_dask_worker_summary(client):
    workers, slots = _dask_worker_slots(client)
    by_host = {}
    for worker in workers.values():
        host = worker.get("host", "unknown")
        by_host.setdefault(host, {"workers": 0, "threads": 0})
        by_host[host]["workers"] += 1
        by_host[host]["threads"] += int(worker.get("nthreads", 1) or 1)
    print(
        "qibotn_dask_workers "
        f"workers={len(workers)} threads={slots} by_host={by_host}",
        flush=True,
    )


def _run_trial_with_debug(fn, *args, **kwargs):
    import os
    import socket

    try:
        from distributed import get_worker

        worker = get_worker()
        worker_address = worker.address
    except Exception:
        worker_address = "unknown"

    method = kwargs.get("method", "unknown")
    pid = os.getpid()
    host = socket.gethostname()
    print(
        "qibotn_trial_start "
        f"worker={worker_address} host={host} pid={pid} method={method}",
        flush=True,
    )
    start = time.perf_counter()
    try:
        trial = fn(*args, **kwargs)
    except Exception as exc:
        elapsed = time.perf_counter() - start
        print(
            "qibotn_trial_error "
            f"worker={worker_address} host={host} pid={pid} "
            f"method={method} seconds={elapsed:.3f} error={exc!r}",
            flush=True,
        )
        raise
    elapsed = time.perf_counter() - start
    print(
        "qibotn_trial_done "
        f"worker={worker_address} host={host} pid={pid} method={method} "
        f"seconds={elapsed:.3f} score={trial.get('score', float('nan')):.6g} "
        f"flops={trial.get('flops', float('nan')):.6g} "
        f"size={trial.get('size', float('nan')):.6g}",
        flush=True,
    )
    return trial


def _patch_cotengra_dask_submit(debug_trials=False):
    global _COTENGRA_DASK_SUBMIT_PATCHED, _DASK_TRIAL_DEBUG
    _DASK_TRIAL_DEBUG = bool(debug_trials)
    if _COTENGRA_DASK_SUBMIT_PATCHED:
        return

    import cotengra.parallel as ctg_parallel
    import cotengra.hyperoptimizers.hyper as hyper

    original_submit = ctg_parallel.submit

    def submit(pool, fn, *args, **kwargs):
        backend = pool.__class__.__module__.split(".", 1)[0]
        if _DASK_TRIAL_DEBUG and backend == "distributed":
            return original_submit(
                pool,
                _run_trial_with_debug,
                fn,
                *args,
                **kwargs,
            )
        return original_submit(pool, fn, *args, **kwargs)

    ctg_parallel.submit = submit
    hyper.submit = submit
    _COTENGRA_DASK_SUBMIT_PATCHED = True


def _patch_cotengra_dask_as_completed():
    """Make cotengra 0.7.5 handle distributed.Future objects.

    This cotengra release routes all parallel futures through
    ``concurrent.futures.as_completed()``, which does not accept dask
    ``distributed.Future`` instances. Keep cotengra's optimizer/reporting logic
    intact and only swap the wait primitive when the futures are from dask.
    """
    global _COTENGRA_DASK_PATCHED
    if _COTENGRA_DASK_PATCHED:
        return

    from cotengra.hyperoptimizers.hyper import HyperOptimizer

    def _get_and_report_next_future(self):
        futures_map = {future: setting for setting, future in self._futures}
        if not futures_map:
            return {
                "score": float("inf"),
                "flops": float("inf"),
                "write": float("inf"),
                "size": float("inf"),
                "time": 0.0,
            }

        future0 = next(iter(futures_map))
        if future0.__class__.__module__.split(".", 1)[0] == "distributed":
            from distributed import as_completed

            deadline = getattr(self, "_qibotn_deadline", None)
            timeout = None if deadline is None else max(0.0, deadline - time.time())
            try:
                future = next(iter(as_completed(futures_map, timeout=timeout)))
            except TimeoutError:
                for future in futures_map:
                    future.cancel()
                self._futures = []
                return {
                    "score": float("inf"),
                    "flops": float("inf"),
                    "write": float("inf"),
                    "size": float("inf"),
                    "time": 0.0,
                }
        else:
            import concurrent.futures as _cf

            future = next(_cf.as_completed(futures_map))

        setting = futures_map[future]
        self._futures = [(s, f) for s, f in self._futures if f is not future]
        try:
            trial = future.result()
        except Exception:
            trial = {
                "score": float("inf"),
                "flops": float("inf"),
                "write": float("inf"),
                "size": float("inf"),
                "time": 0.0,
            }
        self._maybe_report_result(setting, trial)
        return trial

    HyperOptimizer._get_and_report_next_future = _get_and_report_next_future
    _COTENGRA_DASK_PATCHED = True


def _search_chunk(
    tn_bytes,
    output_inds,
    repeats,
    seed,
    max_time,
    slicing_opts,
    optlib=None,
):
    import random, cotengra as ctg

    random.seed(seed)
    tn = pickle.loads(tn_bytes)
    kwargs = {}
    if optlib is not None:
        kwargs["optlib"] = optlib
    opt = ctg.HyperOptimizer(
        methods=SEARCH_METHODS,
        max_repeats=repeats,
        max_time=max_time,
        parallel=False,
        minimize="combo-256",
        slicing_opts=slicing_opts,
        progbar=False,
        **kwargs,
    )
    tree = tn.contraction_tree(optimize=opt, output_inds=output_inds)
    return tree.combo_cost(factor=256), _attach_search_stats(tree, opt)


def _run_single_trial(tn_bytes, output_inds, seed, slicing_opts):
    return _search_chunk(
        tn_bytes,
        output_inds,
        repeats=1,
        seed=seed,
        max_time=None,
        slicing_opts=slicing_opts,
        optlib="random",
    )


def _kill_pool(pool):
    processes = getattr(pool, "_processes", None)
    if processes:
        pids = list(processes.keys())
    else:
        pids = []

    for pid in pids:
        try:
            os.kill(pid, signal.SIGKILL)
        except ProcessLookupError:
            pass
    pool.shutdown(wait=False)


def _serial_search(tn_bytes, output_inds, repeats, seed, max_time, slicing_opts=None, trial_timeout=None):
    import time

    if trial_timeout is None:
        return _search_chunk(
            tn_bytes,
            output_inds,
            repeats=repeats,
            seed=seed,
            max_time=max_time,
            slicing_opts=slicing_opts,
        )

    deadline = time.time() + max_time
    best_cost, best_tree = float("inf"), None

    for i in range(repeats):
        if time.time() >= deadline:
            break
        timeout = min(trial_timeout, deadline - time.time())
        pool = ProcessPoolExecutor(max_workers=1)
        fut = pool.submit(_run_single_trial, tn_bytes, output_inds, seed * 10000 + i, slicing_opts)
        try:
            cost, tree = fut.result(timeout=timeout)
            if cost < best_cost:
                best_cost, best_tree = cost, tree
        except Exception:
            pass
        finally:
            _kill_pool(pool)

    return best_cost, best_tree


def _split_repeats(total_repeats, n_workers):
    n_workers = max(1, int(n_workers))
    total_repeats = max(1, int(total_repeats))
    chunk, extra = divmod(total_repeats, n_workers)
    return [chunk + (1 if i < extra else 0) for i in range(n_workers) if chunk + (1 if i < extra else 0) > 0]


def _processpool_search(tn, output_inds, total_repeats, n_workers, max_time, slicing_opts=None, trial_timeout=None):
    tn_bytes = pickle.dumps(tn)
    repeat_chunks = _split_repeats(total_repeats, n_workers)
    pool = ProcessPoolExecutor(max_workers=len(repeat_chunks))
    futures = []
    for seed, repeats in enumerate(repeat_chunks):
        futures.append(
            pool.submit(
                _serial_search,
                tn_bytes,
                output_inds,
                repeats,
                seed,
                max_time,
                slicing_opts,
                trial_timeout,
            )
        )
    best_cost, best_tree = float("inf"), None
    deadline = time.monotonic() + max_time if max_time is not None else None
    try:
        timeout = None if deadline is None else max(0.0, deadline - time.monotonic())
        for fut in as_completed(futures, timeout=timeout):
            try:
                cost, tree = fut.result()
                if cost < best_cost:
                    best_cost, best_tree = cost, tree
            except Exception:
                pass
    except TimeoutError:
        pass
    finally:
        for fut in futures:
            fut.cancel()
        _kill_pool(pool)
    return best_tree


def _dask_search(
    tn,
    output_inds,
    total_repeats,
    max_time,
    slicing_opts=None,
    dask_address=None,
    n_workers=None,
    optlib=None,
    debug_trials=False,
    close_workers=False,
):
    """Run one centralized cotengra hyper-optimizer over a dask pool.

    With ``dask_address`` this connects to an external distributed scheduler.
    Without it, a local dask cluster is created for single-node smoke testing.
    """
    try:
        from distributed import Client, LocalCluster, get_client
    except ImportError as exc:
        raise ImportError(
            "Dask search requires `distributed`. Install it with "
            "`pip install distributed` or the package extra that provides it."
        ) from exc

    import cotengra as ctg

    _patch_cotengra_dask_as_completed()
    _patch_cotengra_dask_submit(debug_trials=debug_trials)

    close_client = False
    close_cluster = False
    cluster = None

    if dask_address:
        client = Client(dask_address)
        close_client = True
    else:
        try:
            client = get_client()
        except ValueError:
            cluster = LocalCluster(
                n_workers=max(1, int(n_workers or os.cpu_count() or 1)),
                threads_per_worker=1,
                processes=True,
                memory_limit=0,
            )
            client = Client(cluster)
            close_client = True
            close_cluster = True

    kwargs = {}
    if optlib is not None:
        kwargs["optlib"] = optlib

    retire_workers = []
    try:
        workers, worker_slots = _dask_worker_slots(client)
        if close_workers:
            retire_workers = list(workers)
        if debug_trials:
            _print_dask_worker_summary(client)
        if total_repeats < worker_slots:
            print(
                "qibotn_dask_underutilized "
                f"requested_trials={total_repeats} worker_slots={worker_slots} "
                "hint='increase --tn-search-repeats to at least worker_slots'",
                flush=True,
            )
        opt = ctg.HyperOptimizer(
            methods=SEARCH_METHODS,
            max_repeats=total_repeats,
            max_time=max_time,
            parallel=client,
            minimize="combo-256",
            slicing_opts=slicing_opts,
            progbar=False,
            **kwargs,
        )
        opt._num_workers = max(1, worker_slots)
        opt.pre_dispatch = max(1, min(int(total_repeats), worker_slots))
        if max_time is not None:
            opt._qibotn_deadline = time.time() + max_time
        tree = tn.contraction_tree(optimize=opt, output_inds=output_inds)
        return _attach_search_stats(tree, opt)
    finally:
        if close_workers and retire_workers:
            try:
                retired = client.retire_workers(
                    workers=retire_workers,
                    close_workers=True,
                    remove=True,
                )
                print(
                    "qibotn_dask_workers_closed "
                    f"requested={len(retire_workers)} retired={len(retired)}",
                    flush=True,
                )
            except Exception as exc:
                print(
                    "qibotn_dask_workers_close_failed "
                    f"requested={len(retire_workers)} error={exc!r}",
                    flush=True,
                )
        if close_client:
            client.close()
        if close_cluster:
            cluster.close()


def _mpi_search(
    tn,
    output_inds,
    total_repeats,
    max_time,
    n_workers=None,
    slicing_opts=None,
    trial_timeout=None,
    search_backend="processpool",
    dask_address=None,
    debug_trials=False,
    dask_close_workers=False,
):
    comm = MPI.COMM_WORLD
    rank, size = comm.Get_rank(), comm.Get_size()
    search_backend = search_backend or "processpool"

    if search_backend == "dask":
        if not dask_address:
            raise ValueError(
                "MPI + dask search requires an external dask scheduler. Start "
                "dask-scheduler/dask-worker outside mpiexec and pass "
                "`--dask-address tcp://host:8786`."
            )

        payload = None
        if rank == 0:
            try:
                tree = _dask_search(
                    tn,
                    output_inds,
                    total_repeats,
                    max_time,
                    slicing_opts=slicing_opts,
                    dask_address=dask_address,
                    n_workers=n_workers,
                    debug_trials=debug_trials,
                    close_workers=dask_close_workers,
                )
                payload = ("ok", tree)
            except Exception as exc:
                payload = ("error", repr(exc))

        status, value = comm.bcast(payload, root=0)
        if status == "error":
            raise RuntimeError(f"Dask path search failed on rank 0: {value}")
        return value

    repeats_per = max(1, total_repeats // size)

    # Run search work in child processes even when n_workers == 1, so the parent
    # MPI rank can enforce the global timeout by killing active trials.
    local_tree = _processpool_search(
        tn,
        output_inds,
        repeats_per,
        max(1, n_workers or 1),
        max_time,
        slicing_opts,
        trial_timeout,
    )
    local_cost = local_tree.combo_cost(factor=256) if local_tree else float("inf")

    all_results = comm.gather((local_cost, local_tree), root=0)
    best_tree = None
    if rank == 0:
        best_cost = float("inf")
        for cost, tree in all_results:
            if tree is not None and cost < best_cost:
                best_cost, best_tree = cost, tree
    return comm.bcast(best_tree, root=0)


def parallel_path_search(tn, output_inds, method='processpool', total_repeats=1024,
                         max_time=300, n_workers=48, slicing_opts=None,
                         trial_timeout=None, search_backend=None,
                         dask_address=None, debug_trials=False,
                         dask_close_workers=False):
    """Parallel contraction path search.

    Args:
        method: 'processpool' | 'dask' | 'mpi' | 'serial'
        total_repeats: Total optimization repeats across all workers
        max_time: Global timeout per worker (seconds)
        n_workers: Workers per MPI rank (or total for processpool)
        slicing_opts: cotengra slicing options for memory control
        trial_timeout: Per-trial timeout (seconds); kills and skips hung trials
    """
    if method == 'serial':
        tn_bytes = pickle.dumps(tn)
        _, tree = _serial_search(tn_bytes, output_inds, total_repeats, 0, max_time, slicing_opts, trial_timeout)
        return tree
    elif method == 'mpi':
        if not _HAVE_MPI:
            raise ImportError("mpi4py not available")
        return _mpi_search(
            tn,
            output_inds,
            total_repeats,
            max_time,
            n_workers,
            slicing_opts,
            trial_timeout,
            search_backend=search_backend,
            dask_address=dask_address,
            debug_trials=debug_trials,
            dask_close_workers=dask_close_workers,
        )
    elif method == 'processpool':
        return _processpool_search(tn, output_inds, total_repeats, n_workers, max_time, slicing_opts, trial_timeout)
    elif method == 'dask':
        return _dask_search(
            tn,
            output_inds,
            total_repeats,
            max_time,
            slicing_opts=slicing_opts,
            dask_address=dask_address,
            n_workers=n_workers,
            debug_trials=debug_trials,
            close_workers=dask_close_workers,
        )
    else:
        raise ValueError(f"Unknown method: {method}")


def contraction_tree_costs(tree, dtype_bytes=16, combo_factor=256):
    """Return comparable cost estimates for a cotengra contraction tree.

    These values are estimates, not profiling results. They are the right first
    signal for path quality: lower ``combo`` usually means lower CPU contraction
    time, while ``peak_memory_gib`` estimates the largest intermediate tensor.
    """
    stats = tree.contract_stats()
    flops = float(stats["flops"])
    write = float(stats["write"])
    size = float(stats["size"])
    combo = float(tree.combo_cost(factor=combo_factor))
    nslices = int(getattr(tree, "multiplicity", 1))
    original_flops = float(stats.get("original_flops", flops))

    return {
        "flops": flops,
        "write": write,
        "size": size,
        "combo": combo,
        "log10_flops": log10(flops) if flops > 0 else float("-inf"),
        "log10_write": log10(write) if write > 0 else float("-inf"),
        "log2_size": log2(size) if size > 0 else float("-inf"),
        "log10_combo": log10(combo) if combo > 0 else float("-inf"),
        "nslices": nslices,
        "slicing_overhead": flops / original_flops if original_flops > 0 else float("nan"),
        "peak_memory_gib": size * dtype_bytes / 1024**3,
    }


@dataclass(frozen=True)
class SlicePlan:
    """Slice ownership for one MPI rank."""

    rank: int
    size: int
    nslices: int
    indices: tuple
    assignment: str = "block"


@dataclass(frozen=True)
class SlicedContractStats:
    """Diagnostics for one sliced contraction."""

    rank: int
    size: int
    nslices: int
    local_slices: int
    assignment: str


def mpi_slice_plan(nslices, rank, size, assignment="block"):
    """Return the contraction slice ids assigned to one MPI rank.

    ``block`` gives each rank a contiguous range, mirroring cutensornet's
    slice-range style. ``cyclic`` gives rank ``r`` slices ``r, r + size, ...``,
    which can balance better if individual slice costs vary.
    """
    if nslices < 0:
        raise ValueError("nslices must be non-negative.")
    if size <= 0:
        raise ValueError("size must be positive.")
    if not 0 <= rank < size:
        raise ValueError("rank must satisfy 0 <= rank < size.")

    if assignment == "block":
        chunk, extra = divmod(nslices, size)
        start = rank * chunk + min(rank, extra)
        stop = start + chunk + (1 if rank < extra else 0)
        indices = tuple(range(start, stop))
    elif assignment == "cyclic":
        indices = tuple(range(rank, nslices, size))
    else:
        raise ValueError("assignment must be 'block' or 'cyclic'.")

    return SlicePlan(rank, size, nslices, indices, assignment)


def _array_backend(arrays):
    return "torch" if type(arrays[0]).__module__.startswith("torch") else "numpy"


def _to_numpy_vector(value, is_torch):
    if is_torch:
        return value.detach().cpu().numpy().reshape(-1)
    return np.asarray(value).reshape(-1)


def _zero_vector_like(arrays):
    array = arrays[0]
    if type(array).__module__.startswith("torch"):
        return np.zeros(1, dtype=np.complex64 if "64" in str(array.dtype) else np.complex128)
    return np.zeros(1, dtype=np.asarray(array).dtype)


def contract_tree_slices(tree, arrays, slice_indices, backend=None, implementation=None):
    """Contract a subset of cotengra slices and return their local sum."""
    backend = backend or _array_backend(arrays)
    is_torch = backend == "torch"
    local = None
    cpp_contract = None
    if implementation == "cpp":
        if backend != "torch":
            raise ValueError("implementation='cpp' requires torch arrays.")
        from qibotn.torch_contractor import contract_tree_cpp

        cpp_contract = contract_tree_cpp

    for slice_id in slice_indices:
        if cpp_contract is not None:
            value = cpp_contract(tree, tree.slice_arrays(arrays, slice_id))
        elif implementation is None:
            value = tree.contract_slice(arrays, slice_id, backend=backend)
        else:
            value = tree.contract_slice(
                arrays,
                slice_id,
                backend=backend,
                implementation=implementation,
            )
        value = _to_numpy_vector(value, is_torch)
        local = value if local is None else local + value

    return _zero_vector_like(arrays) if local is None else local


def parallel_contract(
    tree,
    arrays,
    method='mpi',
    comm=None,
    assignment="block",
    return_stats=False,
    implementation=None,
):
    if method == 'mpi':
        if not _HAVE_MPI or comm is None:
            raise ValueError("MPI method requires mpi4py and comm")
        return _contract_mpi(
            tree,
            arrays,
            comm,
            assignment=assignment,
            return_stats=return_stats,
            implementation=implementation,
        )
    raise ValueError(f"Unknown method: {method}")


def _contract_mpi(
    tree,
    arrays,
    comm,
    root=0,
    assignment="block",
    return_stats=False,
    implementation=None,
):
    rank, size = comm.Get_rank(), comm.Get_size()
    backend = _array_backend(arrays)
    is_torch = backend == "torch"
    nslices = int(getattr(tree, "multiplicity", 1))
    stats = SlicedContractStats(rank, size, nslices, 0, assignment)
    nslices_by_rank = comm.allgather(nslices)
    if len(set(nslices_by_rank)) != 1:
        raise RuntimeError(
            "Inconsistent contraction tree slices across MPI ranks: "
            f"{nslices_by_rank}. Ensure all nodes load the same tree file."
        )

    if not set(getattr(tree, "sliced_inds", ())).isdisjoint(set(getattr(tree, "output", ()))):
        raise NotImplementedError(
            "MPI sliced contraction currently requires sliced indices not to "
            "appear in the output."
        )

    plan = mpi_slice_plan(nslices, rank, size, assignment=assignment)
    local = contract_tree_slices(
        tree,
        arrays,
        plan.indices,
        backend=backend,
        implementation=implementation,
    )
    stats = SlicedContractStats(rank, size, nslices, len(plan.indices), assignment)

    result = np.zeros_like(local) if rank == root else None
    comm.Reduce(local, result, root=root)
    return (result, stats) if return_stats else result
