"""Parallel path search and contraction utilities for tensor networks."""
import os
import pickle
import signal
import numpy as np
from concurrent.futures import ProcessPoolExecutor, as_completed

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


def _run_single_trial(tn_bytes, output_inds, seed, slicing_opts):
    import random, cotengra as ctg
    random.seed(seed)
    tn = pickle.loads(tn_bytes)
    opt = ctg.HyperOptimizer(
        methods=["kahypar", "kahypar-agglom", "spinglass"],
        max_repeats=1,
        parallel=False,
        minimize="combo-256",
        optlib="random",
        slicing_opts=slicing_opts,
        progbar=False,
    )
    tree = tn.contraction_tree(optimize=opt, output_inds=output_inds)
    return tree.combo_cost(factor=256), tree


def _kill_pool(pool):
    for pid in list(pool._processes.keys()):
        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:
        import random, cotengra as ctg
        random.seed(seed)
        tn = pickle.loads(tn_bytes)
        opt = ctg.HyperOptimizer(
            methods=["kahypar", "kahypar-agglom", "spinglass"],
            max_repeats=repeats,
            parallel=False,
            minimize="combo-256",
            max_time=max_time,
            optlib="random",
            slicing_opts=slicing_opts,
            progbar=False,
        )
        tree = tn.contraction_tree(optimize=opt, output_inds=output_inds)
        return tree.combo_cost(factor=256), tree

    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 _processpool_search(tn, output_inds, total_repeats, n_workers, max_time, slicing_opts=None, trial_timeout=None):
    tn_bytes = pickle.dumps(tn)
    repeats_per = max(1, total_repeats // n_workers)
    pool = ProcessPoolExecutor(max_workers=n_workers)
    futures = [
        pool.submit(_serial_search, tn_bytes, output_inds, repeats_per, seed, max_time, slicing_opts, trial_timeout)
        for seed in range(n_workers)
    ]
    best_cost, best_tree = float("inf"), None
    try:
        for fut in as_completed(futures, timeout=max_time + 5):
            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 _mpi_search(tn, output_inds, total_repeats, max_time, n_workers=None, slicing_opts=None, trial_timeout=None):
    comm = MPI.COMM_WORLD
    rank, size = comm.Get_rank(), comm.Get_size()
    tn_bytes = pickle.dumps(tn)
    repeats_per = max(1, total_repeats // size)

    if n_workers and n_workers > 1:
        local_tree = _processpool_search(
            tn, output_inds, repeats_per, n_workers, max_time, slicing_opts, trial_timeout
        )
        local_cost = local_tree.combo_cost(factor=256) if local_tree else float("inf")
    else:
        local_cost, local_tree = _serial_search(
            tn_bytes, output_inds, repeats_per, rank, max_time, slicing_opts, trial_timeout
        )

    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):
    """Parallel contraction path search.

    Args:
        method: 'processpool' | '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)
    elif method == 'processpool':
        return _processpool_search(tn, output_inds, total_repeats, n_workers, max_time, slicing_opts, trial_timeout)
    else:
        raise ValueError(f"Unknown method: {method}")


def parallel_contract(tree, arrays, method='mpi', comm=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)
    raise ValueError(f"Unknown method: {method}")


def _contract_mpi(tree, arrays, comm, root=0):
    rank, size = comm.Get_rank(), comm.Get_size()
    is_torch = type(arrays[0]).__module__.startswith("torch")

    result_np = None
    for i in range(rank, tree.multiplicity, size):
        x = tree.contract_slice(arrays, i)
        x_np = np.asarray(x.detach().cpu().numpy() if is_torch else x).reshape(-1)
        result_np = x_np if result_np is None else result_np + x_np

    if result_np is None:
        result_np = np.zeros(1, dtype=np.complex128)

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