"""High-level Python interface to hypredrive.

The :class:`HypreDrive` class is a thin orchestration layer over the
private Cython binding. Its job is twofold:

* normalize the various input shapes a user is likely to supply
  (``scipy.sparse.csr_matrix`` vs. raw ``(indptr, indices, data)`` slabs;
  Python ``dict`` of options vs. YAML literal vs. file path);
* enforce a small set of correctness invariants the C API does not
  natively check (e.g. dtypes match the platform's ``HYPRE_BigInt``,
  RHS length matches the matrix row range).

Anything beyond that (the actual solve, error reporting, ownership) is
delegated to the C library. The class is intentionally not a thin
``ctypes`` mirror of every C entry point: that would couple Python users
to internal naming choices and make later refactors painful.
"""

from __future__ import annotations

import contextlib
from typing import Any, Optional, Sequence

import numpy as np

from . import _core, session
from .options import OptionsLike, normalize_options
from .result import SolveResult

__all__ = ["HypreDrive", "solve", "BIGINT_DTYPE", "REAL_DTYPE"]


# ---------------------------------------------------------------------------
# Platform-derived dtypes
#
# HYPRE_BigInt is either ``int32`` or ``int64`` depending on whether HYPRE
# was built with mixed-int support. The exact dtype is observable at native
# build time and exposed through ``_core._hypre_bigint_size``; we surface
# it as a module-level numpy dtype so user code does not have to guess.
# ---------------------------------------------------------------------------

def _resolve_bigint_dtype() -> np.dtype:
    size = _core._hypre_bigint_size()
    if size == 4:
        return np.dtype(np.int32)
    if size == 8:
        return np.dtype(np.int64)
    raise RuntimeError(
        f"Unsupported HYPRE_BigInt size {size} reported by the native module"
    )


def _resolve_real_dtype() -> np.dtype:
    size = _core._hypre_real_size()
    if size == 8:
        return np.dtype(np.float64)
    if size == 4:
        return np.dtype(np.float32)
    raise RuntimeError(
        f"Unsupported HYPRE_Real size {size} reported by the native module"
    )


BIGINT_DTYPE: np.dtype = _resolve_bigint_dtype()
REAL_DTYPE: np.dtype = _resolve_real_dtype()


def _as_bigint_array(arr: Any, name: str) -> np.ndarray:
    a = np.ascontiguousarray(arr, dtype=BIGINT_DTYPE)
    if a.ndim != 1:
        raise ValueError(f"{name} must be a 1-D array, got shape {a.shape}")
    return a


def _as_real_array(arr: Any, name: str) -> np.ndarray:
    a = np.ascontiguousarray(arr, dtype=REAL_DTYPE)
    if a.ndim != 1:
        raise ValueError(f"{name} must be a 1-D array, got shape {a.shape}")
    return a


def _validate_bigint_scalar(value: Any, name: str) -> int:
    ivalue = int(value)
    info = np.iinfo(BIGINT_DTYPE)
    if ivalue < int(info.min) or ivalue > int(info.max):
        raise ValueError(
            f"{name} ({ivalue}) is outside the HYPRE_BigInt range "
            f"[{int(info.min)}, {int(info.max)}]"
        )
    return ivalue


# ---------------------------------------------------------------------------
# HypreDrive: object-oriented entry point
# ---------------------------------------------------------------------------

class HypreDrive:
    """Stateful driver wrapping a single ``HYPREDRV_t`` handle.

    Typical usage::

        with hypredrive.HypreDrive(options=...) as drv:
            drv.set_matrix_from_csr(A_csr)
            drv.set_rhs(b)
            drv.solve()
            x = drv.get_solution()

    The driver always operates in *library mode*: it expects the matrix
    and RHS to be supplied via ``set_matrix_from_csr`` / ``set_rhs``. If
    you want the legacy CLI behavior of reading from a YAML file with
    on-disk matrix/RHS data, use the C driver directly.
    """

    def __init__(
        self,
        options: OptionsLike = None,
        comm: Any = None,
        input_args: Optional[Sequence[str]] = None,
    ) -> None:
        session.initialize()
        self._core: Optional[_core.HypreDriveCore] = _core.HypreDriveCore(
            comm=comm,
            library_mode=True,
        )
        self._row_start: Optional[int] = None
        self._row_end: Optional[int] = None
        self._rhs_set: bool = False
        self._last_iterations: Optional[int] = None
        self._last_converged: Optional[bool] = None
        self._last_final_res_norm: Optional[float] = None
        self._last_setup_time: Optional[float] = None
        self._last_solve_time: Optional[float] = None

        # Always parse a YAML configuration up front so the C side has a
        # solver/preconditioner method selected before we hand over data.
        # ``input_args`` supplies CLI-style overrides applied on top of the
        # options, e.g. ["--solver:pcg:max_iter", "100"] (a leading "-a"
        # token is tolerated).
        yaml_text = normalize_options(options)
        extra = list(input_args) if input_args else None
        self._core.parse_yaml(yaml_text.encode("utf-8"), extra)

    # ------------------------------------------------------------------
    # Context-manager protocol
    # ------------------------------------------------------------------

    def __enter__(self) -> "HypreDrive":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()

    def close(self) -> None:
        """Release the underlying HYPREDRV handle. Idempotent."""
        if self._core is not None:
            self._core.close()
            self._core = None

    # ------------------------------------------------------------------
    # Data ingest
    # ------------------------------------------------------------------

    def set_matrix_from_csr(
        self,
        matrix_or_indptr: Any,
        col_indices: Any = None,
        data: Any = None,
        row_start: Optional[int] = None,
        row_end: Optional[int] = None,
    ) -> None:
        """Install a system matrix from CSR.

        Two call shapes are accepted:

        1. ``set_matrix_from_csr(scipy_csr_matrix)``: the ``indptr``,
           ``indices``, ``data`` triple is taken from the SciPy object.
           ``row_start`` defaults to ``0`` and ``row_end`` to
           ``shape[0] - 1``. If an explicit row range is supplied with a
           SciPy matrix, the matrix is interpreted as this rank's local slab
           and must have ``shape[0] == row_end - row_start + 1``.
        2. ``set_matrix_from_csr(indptr, col_indices, data,
           row_start=..., row_end=...)``: distributed mode, useful when
           each MPI rank owns a slab of rows.
        """
        self._require_open()

        if col_indices is None and data is None:
            if isinstance(matrix_or_indptr, tuple) and len(matrix_or_indptr) == 3:
                indptr, indices, values = matrix_or_indptr
            elif hasattr(matrix_or_indptr, "tocsr"):
                if (row_start is None) != (row_end is None):
                    raise TypeError(
                        "row_start and row_end must be provided together "
                        "when passing a SciPy sparse matrix"
                    )
                if getattr(matrix_or_indptr, "format", None) == "csr":
                    csr = matrix_or_indptr
                else:
                    csr = matrix_or_indptr.tocsr()
                indptr, indices, values = csr.indptr, csr.indices, csr.data
                if row_start is None:
                    row_start = 0
                if row_end is None:
                    row_end = csr.shape[0] - 1
            else:
                raise TypeError(
                    "set_matrix_from_csr expects a scipy.sparse matrix or "
                    "the (indptr, col_indices, data) triple"
                )
        else:
            indptr, indices, values = matrix_or_indptr, col_indices, data

        if row_start is None or row_end is None:
            raise TypeError(
                "row_start and row_end are required when passing raw "
                "(indptr, col_indices, data) arrays"
            )

        indptr_arr = _as_bigint_array(indptr, "indptr")
        indices_arr = _as_bigint_array(indices, "col_indices")
        values_arr = _as_real_array(values, "data")

        row_start_i = _validate_bigint_scalar(row_start, "row_start")
        row_end_i = _validate_bigint_scalar(row_end, "row_end")

        # Keep these checks even though the C builder repeats them. The Python
        # layer gives users immediate shape/range errors, while the C layer
        # remains the trust boundary for direct _core/C API callers.
        nrows = row_end_i - row_start_i + 1
        if nrows < 0:
            raise ValueError(
                f"row_end ({row_end}) must be >= row_start ({row_start})"
            )
        if indptr_arr.shape[0] != nrows + 1:
            raise ValueError(
                f"indptr length {indptr_arr.shape[0]} does not match "
                f"nrows+1 = {nrows + 1}"
            )
        indptr_start = int(indptr_arr[0])
        indptr_end = int(indptr_arr[-1])
        if indptr_start < 0:
            raise ValueError(
                f"indptr[0] ({indptr_start}) must be nonnegative"
            )
        if indptr_end < indptr_start:
            raise ValueError(
                f"indptr[-1] ({indptr_end}) must be >= indptr[0] "
                f"({indptr_start})"
            )
        if (np.diff(indptr_arr) < 0).any():
            raise ValueError("indptr entries must be monotonically nondecreasing")
        required_entries = indptr_end
        if (
            indices_arr.shape[0] < required_entries
            or values_arr.shape[0] < required_entries
        ):
            raise ValueError(
                "col_indices/data must have at least indptr[-1]="
                f"{required_entries} entries"
            )

        # "matrix" marks the start of a new linear system in the stats
        # context (HYPREDRV_AnnotateBegin docs); without it every solve on
        # a reused driver collapses into the first statistics entry.
        with self._annotated("matrix"):
            self._core.set_matrix_from_csr(
                row_start_i, row_end_i, indptr_arr, indices_arr, values_arr,
            )
        self._row_start = row_start_i
        self._row_end = row_end_i

    def set_rhs(
        self,
        values: Any,
        row_start: Optional[int] = None,
        row_end: Optional[int] = None,
    ) -> None:
        """Install the right-hand side vector.

        If the matrix has already been registered via
        :meth:`set_matrix_from_csr`, ``row_start`` / ``row_end`` are
        inferred from it; the caller may still override for clarity.
        """
        self._require_open()
        values_arr = _as_real_array(values, "values")

        if row_start is None:
            row_start = self._row_start
        if row_end is None:
            row_end = self._row_end
        if row_start is None or row_end is None:
            raise RuntimeError(
                "RHS row range unknown; call set_matrix_from_csr first or "
                "pass row_start / row_end explicitly"
            )

        row_start_i = _validate_bigint_scalar(row_start, "row_start")
        row_end_i = _validate_bigint_scalar(row_end, "row_end")
        nrows = row_end_i - row_start_i + 1
        if values_arr.shape[0] != nrows:
            raise ValueError(
                f"RHS length {values_arr.shape[0]} does not match "
                f"row range size {nrows}"
            )

        with self._annotated("rhs"):
            self._core.set_rhs_from_array(row_start_i, row_end_i, values_arr)
        self._rhs_set = True

    def set_dofmap(self, labels: Any) -> None:
        """Install per-row DOF labels for MGR.

        ``labels`` is a 1-D sequence (NumPy array, list, etc.) of small
        non-negative integer codes — one per local row — telling MGR which
        block each row belongs to. The length must match the local row
        range previously registered by :meth:`set_matrix_from_csr`.

        The dofmap is required when the configured preconditioner is MGR
        and is otherwise ignored by the C layer. Calling this method
        replaces any previously installed dofmap.
        """
        self._require_open()
        if self._row_start is None or self._row_end is None:
            raise RuntimeError(
                "set_dofmap(): no matrix set; call set_matrix_from_csr first"
            )
        # Range-check before the dtype cast: np.ascontiguousarray to intc
        # silently wraps int64 values outside the int32 range and treats
        # negative sentinels as valid labels, both of which produce
        # nonsensical MGR partitions far away from this entry point.
        labels_in = np.asarray(labels)
        if labels_in.size:
            labels_range = labels_in
            if not (
                np.issubdtype(labels_in.dtype, np.integer)
                or np.issubdtype(labels_in.dtype, np.bool_)
            ):
                try:
                    labels_float = np.asarray(labels_in, dtype=np.float64)
                except (TypeError, ValueError):
                    raise ValueError("labels must be integer codes") from None
                if not np.isfinite(labels_float).all():
                    raise ValueError("labels must be finite integer codes")
                if not np.equal(labels_float, np.trunc(labels_float)).all():
                    raise ValueError("labels must be integer codes")
                labels_range = labels_float
            info = np.iinfo(np.intc)
            lo = int(labels_range.min())
            hi = int(labels_range.max())
            if lo < 0:
                raise ValueError(
                    f"labels must be non-negative; got minimum {lo}"
                )
            if hi > int(info.max):
                raise ValueError(
                    f"labels max ({hi}) exceeds C int range [0, {int(info.max)}]"
                )
        labels_arr = np.ascontiguousarray(labels_in, dtype=np.intc)
        if labels_arr.ndim != 1:
            raise ValueError(
                f"labels must be a 1-D array, got shape {labels_arr.shape}"
            )
        nrows = self._row_end - self._row_start + 1
        if labels_arr.shape[0] != nrows:
            raise ValueError(
                f"labels length {labels_arr.shape[0]} does not match local "
                f"row count {nrows}"
            )
        with self._annotated("dofmap"):
            self._core.set_dofmap(labels_arr)

    # ------------------------------------------------------------------
    # Solve cycle
    # ------------------------------------------------------------------

    # ------------------------------------------------------------------
    # Stats / Caliper annotations
    # ------------------------------------------------------------------

    def annotate_begin(self, name: str, id: int = -1) -> None:
        """Begin a stats/Caliper annotation region.

        Reserved names ("system"/"matrix", "rhs", "dofmap", "prec",
        "solve", "reset_x0", ...) feed the statistics context directly;
        see the ``HYPREDRV_AnnotateBegin`` C documentation. The driver
        already emits "matrix"/"rhs"/"dofmap"/"reset_x0" from its data
        ingest and solve methods, so manual annotation is only needed for
        custom regions (e.g. Caliper-only "Run-*" spans).
        """
        self._require_open()
        self._core.annotate_begin(name.encode("utf-8"), id)

    def annotate_end(self, name: str, id: int = -1) -> None:
        """End a stats/Caliper annotation region (pairs annotate_begin)."""
        self._require_open()
        self._core.annotate_end(name.encode("utf-8"), id)

    @contextlib.contextmanager
    def _annotated(self, name: str):
        encoded = name.encode("utf-8")
        self._core.annotate_begin(encoded, -1)
        try:
            yield
        finally:
            self._core.annotate_end(encoded, -1)

    def solve(self) -> None:
        """Run setup + apply on the configured solver/preconditioner."""
        self._require_open()
        if self._row_start is None:
            raise RuntimeError("solve(): no matrix set; call set_matrix_from_csr")
        if not self._rhs_set:
            raise RuntimeError("solve(): no RHS set; call set_rhs")

        # Reset before we start so the last_* properties never report stale
        # values from a previous successful solve when this one fails partway.
        self._last_iterations = None
        self._last_converged = None
        self._last_final_res_norm = None
        self._last_setup_time = None
        self._last_solve_time = None

        # Build x0 (controlled by the YAML init_guess_mode), then copy it into
        # the working solution — without the reset the solve always starts
        # from the zeroed working vector regardless of mode. ResetInitialGuess
        # annotates the "reset_x0" stats region internally (wrapping it
        # manually would count the repetition twice).
        self._core.setup_initial_guess()
        self._core.reset_initial_guess()

        self._core.solver_create()
        try:
            self._core.solver_setup()
            self._core.solver_apply()
            self._last_iterations = self._core.solver_iterations()
            # Convergence info is read from the live solver object, so it must
            # be captured before solver_destroy runs in the finally block.
            self._last_converged = self._core.solver_converged()
            self._last_final_res_norm = self._core.solver_final_res_norm()
            self._last_setup_time = self._core.solver_setup_time()
            self._last_solve_time = self._core.solver_solve_time()
        finally:
            # Destroy the solver immediately so a subsequent solve cycle on the
            # same HypreDrive object starts from a clean slate, even if setup or
            # apply failed. The C handles for the linear system itself are retained
            # for inspection.
            self._core.solver_destroy()

    # ------------------------------------------------------------------
    # Result extraction
    # ------------------------------------------------------------------

    @property
    def last_iterations(self) -> Optional[int]:
        """Number of iterations taken by the most recent :meth:`solve` call.

        ``None`` until a successful solve. Reset to ``None`` at the start of
        every :meth:`solve`, so a failed solve always reports ``None`` here
        rather than the previous solve's count.
        """
        return self._last_iterations

    @property
    def last_converged(self) -> Optional[bool]:
        """Whether the most recent :meth:`solve` reached its tolerance.

        ``False`` means the solver stopped without converging (e.g., it hit
        the maximum number of iterations). ``None`` until a successful solve;
        reset to ``None`` at the start of every :meth:`solve`.
        """
        return self._last_converged

    @property
    def last_final_res_norm(self) -> Optional[float]:
        """Final relative residual norm from the most recent :meth:`solve`.

        ``None`` until a successful solve; reset to ``None`` at the start of
        every :meth:`solve`.
        """
        return self._last_final_res_norm

    @property
    def last_setup_time(self) -> Optional[float]:
        """Solver/preconditioner setup time from the most recent :meth:`solve`.

        Unit matches whatever the C library was configured to report:
        seconds by default, milliseconds when the YAML option
        ``general.use_millisec`` is set to ``yes``. Callers that mix the two
        modes must inspect their own YAML to label printouts correctly.
        ``None`` until the first successful solve.
        """
        return self._last_setup_time

    @property
    def last_solve_time(self) -> Optional[float]:
        """Solver-apply (Krylov iteration) time from the most recent :meth:`solve`.

        Returned in the same units as :attr:`last_setup_time` — see that
        property for the unit/YAML coupling. ``None`` until the first
        successful solve.
        """
        return self._last_solve_time

    def get_solution(self) -> np.ndarray:
        """Return the local-rank solution slab as a NumPy array."""
        self._require_open()
        if self._row_start is None or self._row_end is None:
            raise RuntimeError("get_solution(): solve() has not been run yet")
        nrows = self._row_end - self._row_start + 1
        out = np.empty(nrows, dtype=REAL_DTYPE)
        self._core.copy_solution(out)
        return out

    def solution_norm(self, kind: str = "l2") -> float:
        """Return the norm of the current solution as computed by HYPRE."""
        self._require_open()
        if kind not in {"l1", "l2", "linf"}:
            raise ValueError("kind must be one of 'l1', 'l2', or 'linf'")
        return self._core.solution_norm(kind)

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    def _require_open(self) -> None:
        if self._core is None:
            raise RuntimeError("HypreDrive is closed")


# ---------------------------------------------------------------------------
# One-shot solve helper
# ---------------------------------------------------------------------------

def solve(
    A: Any,
    b: Any,
    options: OptionsLike = None,
    comm: Any = None,
    *,
    row_start: Optional[int] = None,
    row_end: Optional[int] = None,
    input_args: Optional[Sequence[str]] = None,
) -> SolveResult:
    """Configure, solve, and tear down a single linear system.

    This is sugar over the :class:`HypreDrive` lifecycle for the common
    case where the caller does not need to reuse the driver across solves.

    Parameters
    ----------
    A:
        ``scipy.sparse.csr_matrix`` (or anything with a ``.tocsr()``
        method) on a single rank, or any object accepted by
        :meth:`HypreDrive.set_matrix_from_csr`.
    b:
        Local-rank RHS values.
    options:
        Solver/preconditioner configuration. See
        :func:`hypredrive.options.normalize_options` for accepted shapes.
    comm:
        Optional ``mpi4py.MPI.Comm``. Defaults to ``MPI_COMM_SELF``.
    row_start, row_end:
        Inclusive local row range, required when ``A`` is not a SciPy
        matrix.
    input_args:
        CLI-style override tokens applied on top of ``options``, e.g.
        ``["--solver:pcg:max_iter", "100"]``.
    """
    if (row_start is None) != (row_end is None):
        raise TypeError("row_start and row_end must be provided together")

    with HypreDrive(options=options, comm=comm, input_args=input_args) as drv:
        if row_start is None or row_end is None:
            drv.set_matrix_from_csr(A)
        else:
            drv.set_matrix_from_csr(A, row_start=row_start, row_end=row_end)
        drv.set_rhs(b, row_start=row_start, row_end=row_end)
        drv.solve()
        x = drv.get_solution()
        norm = drv.solution_norm("l2")
    return SolveResult(x=x, solution_norm=norm)
