#   Copyright 2024 - present The PyMC Developers
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

"""Storage backends for traces.

The NDArray (pymc.backends.NDArray) backend holds the entire trace in memory.

Selecting values from a backend
-------------------------------

After a backend is finished sampling, it returns a MultiTrace object.
Values can be accessed in a few ways. The easiest way is to index the
backend object with a variable or variable name.

    >>> trace["x"]  # or trace.x or trace[x]

The call will return the sampling values of `x`, with the values for
all chains concatenated. (For a single call to `sample`, the number of
chains will correspond to the `cores` argument.)

To discard the first N values of each chain, slicing syntax can be
used.

    >>> trace["x", 1000:]

The `get_values` method offers more control over which values are
returned. The call below will discard the first 1000 iterations
from each chain and keep the values for each chain as separate arrays.

    >>> trace.get_values("x", burn=1000, combine=False)

The `chains` parameter of `get_values` can be used to limit the chains
that are retrieved.

    >>> trace.get_values("x", burn=1000, chains=[0, 2])

MultiTrace objects also support slicing. For example, the following
call would return a new trace object without the first 1000 sampling
iterations for all traces and variables.

    >>> sliced_trace = trace[1000:]

The backend for the new trace is always NDArray, regardless of the
type of original trace.

Loading a saved backend
-----------------------

Saved backends can be loaded using `arviz.from_netcdf`

"""

from __future__ import annotations

import importlib.util
import warnings

from collections.abc import Mapping, Sequence
from copy import copy
from typing import TYPE_CHECKING, Optional, TypeAlias, Union

import numpy as np

from pytensor.tensor.variable import TensorVariable

from pymc.backends.arviz import predictions_to_inference_data, to_inference_data
from pymc.backends.base import BaseTrace, IBaseTrace
from pymc.backends.ndarray import NDArray
from pymc.blocking import PointType
from pymc.model import Model
from pymc.step_methods.compound import BlockedStep, CompoundStep

if TYPE_CHECKING:
    from pymc.backends.zarr import ZarrTrace


class _ZarrChainBase:
    """Marker base for :class:`pymc.backends.zarr.ZarrChain`.

    Importing the real class pulls the optional ``zarr`` dependency. Callers that
    only need an ``isinstance`` check (e.g. dispatch in ``pymc.sampling``) can
    import this base instead and skip the zarr import. The method stubs mirror
    the subset of ``ZarrChain`` invoked through this base so type checkers can
    validate calls guarded by ``isinstance(_, _ZarrChainBase)``.
    """

    def link_stepper(self, step_method) -> None: ...
    def record_sampling_state(self, step=None) -> None: ...
    def store_sampling_state(self, sampling_state) -> None: ...


class _ZarrTraceBase:
    """Marker base for :class:`pymc.backends.zarr.ZarrTrace`. See ``_ZarrChainBase``."""

    straces: Sequence[BaseTrace]

    def init_trace(self, *args, **kwargs) -> None: ...
    def split_warmup_groups(self) -> None: ...


HAS_MCB = importlib.util.find_spec("mcbackend") is not None

if TYPE_CHECKING and HAS_MCB:
    from mcbackend import Backend, Run

    TraceOrBackend: TypeAlias = BaseTrace | Backend
    RunType: TypeAlias = Run
else:
    TraceOrBackend = BaseTrace  # type: ignore[assignment, misc]
    RunType = type(None)  # type: ignore[assignment, misc]


__all__ = ["predictions_to_inference_data", "to_inference_data"]


# ZarrTrace requires the optional zarr dependency, so we keep it lazy.
def __getattr__(name):
    if name == "ZarrTrace":
        from pymc.backends.zarr import ZarrTrace

        return ZarrTrace
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")


def _init_trace(
    *,
    expected_length: int,
    chain_number: int,
    stats_dtypes: list[dict[str, type]],
    trace: BaseTrace | None,
    model: Model,
    trace_vars: list[TensorVariable] | None = None,
    initial_point: PointType | None = None,
) -> BaseTrace:
    """Initialize a trace backend for a chain."""
    strace: BaseTrace
    if trace is None:
        strace = NDArray(model=model, vars=trace_vars, test_point=initial_point)
    elif isinstance(trace, BaseTrace):
        if len(trace) > 0:
            raise ValueError("Continuation of traces is no longer supported.")
        strace = copy(trace)
    else:
        raise NotImplementedError(f"Unsupported `trace`: {trace}")

    strace.setup(expected_length, chain_number, stats_dtypes)
    return strace


def init_traces(
    *,
    backend: TraceOrBackend | ZarrTrace | None,
    chains: int,
    expected_length: int,
    step: BlockedStep | CompoundStep,
    initial_point: PointType,
    model: Model,
    trace_vars: list[TensorVariable] | None = None,
    tune: int = 0,
) -> tuple[RunType | None, Sequence[IBaseTrace]]:
    """Initialize a trace recorder for each chain."""
    if isinstance(backend, _ZarrTraceBase):
        backend.init_trace(
            chains=chains,
            draws=expected_length - tune,
            tune=tune,
            step=step,
            model=model,
            vars=trace_vars,
            test_point=initial_point,
        )
        return None, backend.straces
    if HAS_MCB:
        try:
            with warnings.catch_warnings():
                # mcbackend touches arviz.InferenceData on import; the
                # MigrationWarning is emitted from arviz, so filter by message.
                warnings.filterwarnings("ignore", "arviz.InferenceData")
                from mcbackend import Backend
        except ImportError:
            pass
        else:
            if isinstance(backend, Backend):
                from pymc.backends.mcbackend import init_chain_adapters

                return init_chain_adapters(
                    backend=backend,
                    chains=chains,
                    initial_point=initial_point,
                    step=step,
                    model=model,
                )

    assert backend is None or isinstance(backend, BaseTrace)
    traces = [
        _init_trace(
            expected_length=expected_length,
            stats_dtypes=step.stats_dtypes,
            chain_number=chain_number,
            trace=backend,
            model=model,
            trace_vars=trace_vars,
            initial_point=initial_point,
        )
        for chain_number in range(chains)
    ]
    return None, traces
