#   Copyright 2022 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.

"""Utilities for converting Pathfinder results to xarray and adding them to InferenceData."""

from __future__ import annotations

import warnings

from dataclasses import asdict

import arviz as az
import numpy as np
import pymc as pm
import xarray as xr

from pymc.blocking import DictToArrayBijection

from pymc_extras.inference.pathfinder.lbfgs import LBFGSStatus
from pymc_extras.inference.pathfinder.pathfinder import (
    MultiPathfinderResult,
    PathfinderConfig,
    PathfinderResult,
    PathStatus,
)


def get_param_coords(model: pm.Model | None, n_params: int) -> list[str]:
    """
    Get parameter coordinate labels from PyMC model.

    Parameters
    ----------
    model : pm.Model | None
        PyMC model to extract variable names from. If None, returns numeric indices.
    n_params : int
        Number of parameters (for fallback indexing when model is None)

    Returns
    -------
    list[str]
        Parameter coordinate labels
    """
    if model is None:
        return [str(i) for i in range(n_params)]

    ip = model.initial_point()
    bij = DictToArrayBijection.map(ip)

    coords = []
    for var_name, shape, size, _ in bij.point_map_info:
        if size == 1:
            coords.append(var_name)
        else:
            for i in range(size):
                coords.append(f"{var_name}[{i}]")
    return coords


def _status_counter_to_dataarray(counter, status_enum_cls) -> xr.DataArray:
    """Convert a Counter of status values to a dense xarray DataArray."""
    all_statuses = list(status_enum_cls)
    status_names = [s.name for s in all_statuses]

    counts = np.array([counter.get(status, 0) for status in all_statuses])

    return xr.DataArray(
        counts, dims=["status"], coords={"status": status_names}, name="status_counts"
    )


def _extract_scalar(value):
    """Extract scalar from array-like or return as-is."""
    if hasattr(value, "item"):
        return value.item()
    elif hasattr(value, "__len__") and len(value) == 1:
        return value[0]
    return value


def pathfinder_result_to_xarray(
    result: PathfinderResult,
    model: pm.Model | None = None,
) -> xr.Dataset:
    """
    Convert a PathfinderResult to an xarray Dataset.

    Parameters
    ----------
    result : PathfinderResult
        Single pathfinder run result
    model : pm.Model | None
        PyMC model for parameter name extraction

    Returns
    -------
    xr.Dataset
        Dataset with pathfinder results

    Examples
    --------
    >>> import pymc as pm
    >>> import pymc_extras as pmx
    >>>
    >>> with pm.Model() as model:
    ...     x = pm.Normal("x", 0, 1)
    ...     y = pm.Normal("y", x, 1, observed=2.0)
    >>> # Assuming we have a PathfinderResult from a pathfinder run
    >>> ds = pathfinder_result_to_xarray(result, model=model)
    >>> print(ds.data_vars)  # Shows lbfgs_niter, elbo_argmax, status info, etc.
    >>> print(ds.attrs)  # Shows metadata like lbfgs_status, path_status
    """
    data_vars = {}
    coords = {}
    attrs = {}

    n_params = None
    if result.samples is not None:
        n_params = result.samples.shape[-1]
    elif hasattr(result, "lbfgs_niter") and result.lbfgs_niter is not None:
        if model is not None:
            try:
                ip = model.initial_point()
                n_params = len(DictToArrayBijection.map(ip).data)
            except Exception:
                pass

    if n_params is not None:
        coords["param"] = get_param_coords(model, n_params)

    if result.lbfgs_niter is not None:
        data_vars["lbfgs_niter"] = xr.DataArray(_extract_scalar(result.lbfgs_niter))

    if result.elbo_argmax is not None:
        data_vars["elbo_argmax"] = xr.DataArray(_extract_scalar(result.elbo_argmax))

    data_vars["lbfgs_status_code"] = xr.DataArray(result.lbfgs_status.value)
    data_vars["lbfgs_status_name"] = xr.DataArray(result.lbfgs_status.name)
    data_vars["path_status_code"] = xr.DataArray(result.path_status.value)
    data_vars["path_status_name"] = xr.DataArray(result.path_status.name)

    if n_params is not None and result.samples is not None:
        if result.samples.ndim >= 2:
            representative_sample = result.samples[0, -1, :]
            data_vars["final_sample"] = xr.DataArray(
                representative_sample, dims=["param"], coords={"param": coords["param"]}
            )

    if result.logP is not None:
        logP = result.logP.flatten() if hasattr(result.logP, "flatten") else result.logP
        if hasattr(logP, "__len__") and len(logP) > 0:
            data_vars["logP_mean"] = xr.DataArray(np.mean(logP))
            data_vars["logP_std"] = xr.DataArray(np.std(logP))
            data_vars["logP_max"] = xr.DataArray(np.max(logP))

    if result.logQ is not None:
        logQ = result.logQ.flatten() if hasattr(result.logQ, "flatten") else result.logQ
        if hasattr(logQ, "__len__") and len(logQ) > 0:
            data_vars["logQ_mean"] = xr.DataArray(np.mean(logQ))
            data_vars["logQ_std"] = xr.DataArray(np.std(logQ))
            data_vars["logQ_max"] = xr.DataArray(np.max(logQ))

    attrs["lbfgs_status"] = result.lbfgs_status.name
    attrs["path_status"] = result.path_status.name

    ds = xr.Dataset(data_vars, coords=coords, attrs=attrs)

    return ds


def multipathfinder_result_to_xarray(
    result: MultiPathfinderResult,
    model: pm.Model | None = None,
    *,
    store_diagnostics: bool = False,
) -> xr.Dataset:
    """
    Convert a MultiPathfinderResult to a single consolidated xarray Dataset.

    Parameters
    ----------
    result : MultiPathfinderResult
        Multi-path pathfinder result
    model : pm.Model | None
        PyMC model for parameter name extraction
    store_diagnostics : bool
        Whether to include potentially large diagnostic arrays

    Returns
    -------
    xr.Dataset
        Single consolidated dataset with all pathfinder results

    Examples
    --------
    >>> import pymc as pm
    >>> import pymc_extras as pmx
    >>>
    >>> with pm.Model() as model:
        ...     x = pm.Normal("x", 0, 1)
    ...
    >>> # Assuming we have a MultiPathfinderResult from multiple pathfinder runs
    >>> ds = multipathfinder_result_to_xarray(result, model=model)
    >>> print("All data:", ds.data_vars)
    >>> print(
    ...     "Summary:",
    ...     [
    ...         k
    ...         for k in ds.data_vars.keys()
    ...         if not k.startswith(("paths/", "config/", "diagnostics/"))
    ...     ],
    ... )
    >>> print("Per-path:", [k for k in ds.data_vars.keys() if k.startswith("paths/")])
    >>> print("Config:", [k for k in ds.data_vars.keys() if k.startswith("config/")])
    """
    n_params = result.samples.shape[-1] if result.samples is not None else None
    param_coords = get_param_coords(model, n_params) if n_params is not None else None

    data_vars = {}
    coords = {}
    attrs = {}

    # Add parameter coordinates if available
    if param_coords is not None:
        coords["param"] = param_coords

    # Build summary-level data (top level)
    _add_summary_data(result, data_vars, coords, attrs)

    # Build per-path data (with paths/ prefix)
    if not result.all_paths_failed and result.samples is not None:
        _add_paths_data(result, data_vars, coords, param_coords, n_params)

    # Build configuration data (with config/ prefix)
    if result.pathfinder_config is not None:
        _add_config_data(result.pathfinder_config, data_vars)

    # Build diagnostics data (with diagnostics/ prefix) if requested
    if store_diagnostics:
        _add_diagnostics_data(result, data_vars, coords, param_coords)

    return xr.Dataset(data_vars, coords=coords, attrs=attrs)


def _add_summary_data(
    result: MultiPathfinderResult, data_vars: dict, coords: dict, attrs: dict
) -> None:
    """Add summary-level statistics to the pathfinder dataset."""
    if result.num_paths is not None:
        data_vars["num_paths"] = xr.DataArray(result.num_paths)
    if result.num_draws is not None:
        data_vars["num_draws"] = xr.DataArray(result.num_draws)

    if result.compile_time is not None:
        data_vars["compile_time"] = xr.DataArray(result.compile_time)
    if result.compute_time is not None:
        data_vars["compute_time"] = xr.DataArray(result.compute_time)
        if result.compile_time is not None:
            data_vars["total_time"] = xr.DataArray(result.compile_time + result.compute_time)

    data_vars["importance_sampling_method"] = xr.DataArray(result.importance_sampling or "none")
    if result.pareto_k is not None:
        data_vars["pareto_k"] = xr.DataArray(result.pareto_k)

    if result.lbfgs_status:
        data_vars["lbfgs_status_counts"] = _status_counter_to_dataarray(
            result.lbfgs_status, LBFGSStatus
        )
    if result.path_status:
        data_vars["path_status_counts"] = _status_counter_to_dataarray(
            result.path_status, PathStatus
        )

    data_vars["all_paths_failed"] = xr.DataArray(result.all_paths_failed)
    if not result.all_paths_failed and result.samples is not None:
        data_vars["num_successful_paths"] = xr.DataArray(result.samples.shape[0])

    if result.lbfgs_niter is not None:
        data_vars["lbfgs_niter_mean"] = xr.DataArray(np.mean(result.lbfgs_niter))
        data_vars["lbfgs_niter_std"] = xr.DataArray(np.std(result.lbfgs_niter))

    if result.elbo_argmax is not None:
        data_vars["elbo_argmax_mean"] = xr.DataArray(np.mean(result.elbo_argmax))
        data_vars["elbo_argmax_std"] = xr.DataArray(np.std(result.elbo_argmax))

    if result.logP is not None:
        data_vars["logP_mean"] = xr.DataArray(np.mean(result.logP))
        data_vars["logP_std"] = xr.DataArray(np.std(result.logP))
        data_vars["logP_max"] = xr.DataArray(np.max(result.logP))

    if result.logQ is not None:
        data_vars["logQ_mean"] = xr.DataArray(np.mean(result.logQ))
        data_vars["logQ_std"] = xr.DataArray(np.std(result.logQ))
        data_vars["logQ_max"] = xr.DataArray(np.max(result.logQ))

    # Add warnings to attributes
    if result.warnings:
        attrs["warnings"] = list(result.warnings)


def _add_paths_data(
    result: MultiPathfinderResult,
    data_vars: dict,
    coords: dict,
    param_coords: list[str] | None,
    n_params: int | None,
) -> None:
    """Add per-path diagnostics to the pathfinder dataset with 'paths/' prefix."""
    n_paths = _determine_num_paths(result)

    # Add path coordinate
    coords["path"] = list(range(n_paths))

    def _add_path_scalar(name: str, data):
        """Add a per-path scalar array to data_vars with paths/ prefix."""
        if data is not None:
            data_vars[f"paths/{name}"] = xr.DataArray(
                data, dims=["path"], coords={"path": coords["path"]}
            )

    _add_path_scalar("lbfgs_niter", result.lbfgs_niter)
    _add_path_scalar("elbo_argmax", result.elbo_argmax)

    if result.logP is not None:
        _add_path_scalar("logP_mean", np.mean(result.logP, axis=1))
        _add_path_scalar("logP_max", np.max(result.logP, axis=1))

    if result.logQ is not None:
        _add_path_scalar("logQ_mean", np.mean(result.logQ, axis=1))
        _add_path_scalar("logQ_max", np.max(result.logQ, axis=1))

    if n_params is not None and result.samples is not None and result.samples.ndim >= 3:
        final_samples = result.samples[:, -1, :]  # (S, N)
        data_vars["paths/final_sample"] = xr.DataArray(
            final_samples,
            dims=["path", "param"],
            coords={"path": coords["path"], "param": coords["param"]},
        )


def _add_config_data(config: PathfinderConfig, data_vars: dict) -> None:
    """Add configuration parameters to the pathfinder dataset with 'config/' prefix."""
    config_dict = asdict(config)
    for key, value in config_dict.items():
        data_vars[f"config/{key}"] = xr.DataArray(value)


def _add_diagnostics_data(
    result: MultiPathfinderResult, data_vars: dict, coords: dict, param_coords: list[str] | None
) -> None:
    """Add detailed diagnostics to the pathfinder dataset with 'diagnostics/' prefix."""
    if result.logP is not None:
        n_paths, n_draws_per_path = result.logP.shape
        if "path" not in coords:
            coords["path"] = list(range(n_paths))
        coords["draw_per_path"] = list(range(n_draws_per_path))

        data_vars["diagnostics/logP_full"] = xr.DataArray(
            result.logP,
            dims=["path", "draw_per_path"],
            coords={"path": coords["path"], "draw_per_path": coords["draw_per_path"]},
        )

    if result.logQ is not None:
        if "draw_per_path" not in coords:
            n_paths, n_draws_per_path = result.logQ.shape
            if "path" not in coords:
                coords["path"] = list(range(n_paths))
            coords["draw_per_path"] = list(range(n_draws_per_path))

        data_vars["diagnostics/logQ_full"] = xr.DataArray(
            result.logQ,
            dims=["path", "draw_per_path"],
            coords={"path": coords["path"], "draw_per_path": coords["draw_per_path"]},
        )

    if result.samples is not None and result.samples.ndim == 3 and param_coords is not None:
        n_paths, n_draws_per_path, n_params = result.samples.shape

        if "path" not in coords:
            coords["path"] = list(range(n_paths))
        if "draw_per_path" not in coords:
            coords["draw_per_path"] = list(range(n_draws_per_path))

        data_vars["diagnostics/samples_full"] = xr.DataArray(
            result.samples,
            dims=["path", "draw_per_path", "param"],
            coords={
                "path": coords["path"],
                "draw_per_path": coords["draw_per_path"],
                "param": coords["param"],
            },
        )


def _determine_num_paths(result: MultiPathfinderResult) -> int:
    """
    Determine the number of paths from per-path arrays.

    When importance sampling is applied, result.samples may be collapsed,
    so we use per-path diagnostic arrays to determine the true path count.
    """
    if result.lbfgs_niter is not None:
        return len(result.lbfgs_niter)
    elif result.elbo_argmax is not None:
        return len(result.elbo_argmax)
    elif result.logP is not None:
        return result.logP.shape[0]
    elif result.logQ is not None:
        return result.logQ.shape[0]

    if result.lbfgs_status:
        return sum(result.lbfgs_status.values())
    elif result.path_status:
        return sum(result.path_status.values())

    if result.samples is not None:
        return result.samples.shape[0]

    raise ValueError("Cannot determine number of paths from result")


def add_pathfinder_to_inference_data(
    idata: az.InferenceData,
    result: PathfinderResult | MultiPathfinderResult,
    model: pm.Model | None = None,
    *,
    group: str = "pathfinder",
    paths_group: str = "pathfinder_paths",  # Deprecated, kept for API compatibility
    diagnostics_group: str = "pathfinder_diagnostics",  # Deprecated, kept for API compatibility
    config_group: str = "pathfinder_config",  # Deprecated, kept for API compatibility
    store_diagnostics: bool = False,
) -> az.InferenceData:
    """
    Add pathfinder results to an ArviZ InferenceData object as a single consolidated group.

    All pathfinder output is now consolidated under a single group with nested structure:
    - Summary statistics at the top level
    - Per-path data with 'paths/' prefix
    - Configuration with 'config/' prefix
    - Diagnostics with 'diagnostics/' prefix (if store_diagnostics=True)

    Parameters
    ----------
    idata : az.InferenceData
        InferenceData object to modify
    result : PathfinderResult | MultiPathfinderResult
        Pathfinder results to add
    model : pm.Model | None
        PyMC model for parameter name extraction
    group : str
        Name for the pathfinder group (default: "pathfinder")
    paths_group : str
        Deprecated: no longer used, kept for API compatibility
    diagnostics_group : str
        Deprecated: no longer used, kept for API compatibility
    config_group : str
        Deprecated: no longer used, kept for API compatibility
    store_diagnostics : bool
        Whether to include potentially large diagnostic arrays

    Returns
    -------
    az.InferenceData
        Modified InferenceData object with consolidated pathfinder group added

    Examples
    --------
    >>> import pymc as pm
    >>> import pymc_extras as pmx
    >>>
    >>> with pm.Model() as model:
    ...     x = pm.Normal("x", 0, 1)
    ...     idata = pmx.fit(method="pathfinder", model=model, add_pathfinder_groups=False)
    >>> # Assuming we have pathfinder results
    >>> idata = add_pathfinder_to_inference_data(idata, results, model=model)
    >>> print(list(idata.groups()))  # Will show ['posterior', 'pathfinder']
    >>> # Access nested data:
    >>> print(
    ...     [k for k in idata.pathfinder.data_vars.keys() if k.startswith("paths/")]
    ... )  # Per-path data
    >>> print(
    ...     [k for k in idata.pathfinder.data_vars.keys() if k.startswith("config/")]
    ... )  # Config data
    """
    # Detect if this is a multi-path result
    # Use isinstance() as primary check, but fall back to duck typing for compatibility
    # with mocks and testing (MultiPathfinderResult has Counter-type status fields)
    is_multipath = isinstance(result, MultiPathfinderResult) or (
        hasattr(result, "lbfgs_status")
        and hasattr(result.lbfgs_status, "values")
        and callable(getattr(result.lbfgs_status, "values"))
    )

    if is_multipath:
        consolidated_ds = multipathfinder_result_to_xarray(
            result, model=model, store_diagnostics=store_diagnostics
        )
    else:
        consolidated_ds = pathfinder_result_to_xarray(result, model=model)

    if group in idata.groups():
        warnings.warn(f"Group '{group}' already exists in InferenceData, it will be replaced.")

    idata.add_groups({group: consolidated_ds})
    return idata
