from getdist import loadMCSamples
from typing import List, Optional
import yaml
from pathlib import Path
from rich.table import Table
from rich.panel import Panel
from .console import console
from .io import print_warning, print_success, create_chain_tree


def get_autocorr(chains: dict) -> None:
    """
    Compute and store autocorrelation information for each MCMC chain.

    .. deprecated:: 0.0.4
        Use :func:`inspect_chains` instead for more comprehensive statistics.

    Args:
        chains: Dictionary of MCSamples objects with labels as keys
    """
    autocorr_table = Table(title="Chain Autocorrelation Statistics", show_header=True)
    autocorr_table.add_column("Chain", style="cyan")
    autocorr_table.add_column("Total Samples", justify="right")
    autocorr_table.add_column("Effective Samples", justify="right", style="yellow")
    autocorr_table.add_column("Autocorr Factor", justify="right")
    autocorr_table.add_column("Efficiency", justify="right", style="green")
    autocorr_table.add_column("Eff/Param", justify="right", style="dim")

    for label, chain in chains.items():
        try:
            n_eff = chain.getEffectiveSamples()
            n_samples = chain.numrows
            autocorr_reduction = n_samples / n_eff
            efficiency = n_eff / n_samples
            n_params = chain.getParamNames().numParams()
            eff_per_param = n_eff / n_params

            autocorr_table.add_row(
                label,
                f"{n_samples:,}",
                f"{n_eff:,.0f}",
                f"{autocorr_reduction:.1f}",
                f"{efficiency:.1%}",
                f"{eff_per_param:.1f}",
            )
        except Exception:
            print_warning(f"Could not compute autocorrelation for chain '{label}'")

    console.print(autocorr_table)


def inspect_chains(
    chains: dict,
    show_params: bool = False,
    show_ranges: bool = False,
    max_params_display: int = 10,
) -> None:
    """
    Display comprehensive statistics and diagnostics for MCMC chains.

    Shows:
        - Basic chain information (samples, parameters, burn-in)
        - Gelman-Rubin R-1 convergence diagnostic
        - Autocorrelation statistics
        - Optionally: parameter names and ranges

    Args:
        chains: Dictionary of MCSamples objects with labels as keys
        show_params: Whether to display parameter names
        show_ranges: Whether to display parameter ranges/bounds
        max_params_display: Maximum number of parameters to display if show_params=True

    Example:
        >>> import primefeat as pf
        >>> chains = pf.get_chains('1e-3')
        >>> pf.inspect_chains(chains)
        >>> pf.inspect_chains(chains, show_params=True, show_ranges=True)
    """
    # Main statistics table
    stats_table = Table(title="Chain Statistics Summary", show_header=True, expand=True)
    stats_table.add_column("Chain", style="cyan", width=15)
    stats_table.add_column("Samples", justify="right", width=10)
    stats_table.add_column("Parameters", justify="right", width=10)
    stats_table.add_column("R-1", justify="right", style="yellow", width=8)
    stats_table.add_column("Status", style="dim", width=15)
    stats_table.add_column("Eff. Samples", justify="right", style="green", width=12)
    stats_table.add_column("Efficiency", justify="right", width=10)

    for label, chain in chains.items():
        try:
            # Basic info
            n_samples = chain.numrows
            n_params = chain.getParamNames().numParams()

            # Gelman-Rubin
            try:
                r_minus_1 = chain.getGelmanRubin()
                if r_minus_1 < 0.01:
                    status = "[green]✓ Converged[/green]"
                    r1_style = "[green]"
                elif r_minus_1 < 0.05:
                    status = "[yellow]~ Acceptable[/yellow]"
                    r1_style = "[yellow]"
                else:
                    status = "[red]✗ Not converged[/red]"
                    r1_style = "[red]"
                r1_str = f"{r1_style}{r_minus_1:.4f}[/{r1_style.strip('[')}]"
            except:
                r_minus_1 = None
                status = "[dim]N/A[/dim]"
                r1_str = "[dim]N/A[/dim]"

            # Autocorrelation
            try:
                n_eff = chain.getEffectiveSamples()
                efficiency = n_eff / n_samples
                n_eff_str = f"{n_eff:,.0f}"
                efficiency_str = f"{efficiency:.1%}"
            except:
                n_eff_str = "[dim]N/A[/dim]"
                efficiency_str = "[dim]N/A[/dim]"

            stats_table.add_row(
                label,
                f"{n_samples:,}",
                str(n_params),
                r1_str,
                status,
                n_eff_str,
                efficiency_str,
            )
        except Exception as e:
            print_warning(f"Could not compute statistics for chain '{label}': {e}")

    console.print(stats_table)

    # Detailed autocorrelation table
    console.print()
    autocorr_table = Table(title="Autocorrelation Details", show_header=True)
    autocorr_table.add_column("Chain", style="cyan")
    autocorr_table.add_column("Autocorr Factor", justify="right", style="yellow")
    autocorr_table.add_column("Eff/Param", justify="right", style="green")
    autocorr_table.add_column("Eff/Sample", justify="right", style="dim")

    for label, chain in chains.items():
        try:
            n_eff = chain.getEffectiveSamples()
            n_samples = chain.numrows
            n_params = chain.getParamNames().numParams()
            autocorr_factor = n_samples / n_eff
            eff_per_param = n_eff / n_params
            eff_per_sample = n_eff / n_samples

            autocorr_table.add_row(
                label,
                f"{autocorr_factor:.2f}",
                f"{eff_per_param:,.1f}",
                f"{eff_per_sample:.3f}",
            )
        except Exception:
            pass

    console.print(autocorr_table)

    # Optional: Parameter information
    if show_params:
        console.print()
        for label, chain in chains.items():
            param_names = chain.getParamNames().names
            n_params = len(param_names)

            # Limit display
            display_params = param_names[:max_params_display]
            has_more = n_params > max_params_display

            param_table = Table(
                title=f"Parameters: {label}",
                show_header=True,
                width=80,
            )
            param_table.add_column("Name", style="cyan")

            if show_ranges:
                param_table.add_column("Mean", justify="right", style="yellow")
                param_table.add_column("Std Dev", justify="right", style="green")
                param_table.add_column("Range", justify="right", style="dim")

            for param_name in display_params:
                try:
                    if show_ranges:
                        samples = chain.getParams().__dict__[param_name]
                        mean = samples.mean()
                        std = samples.std()
                        pmin = samples.min()
                        pmax = samples.max()

                        param_table.add_row(
                            param_name,
                            f"{mean:.4g}",
                            f"{std:.4g}",
                            f"[{pmin:.4g}, {pmax:.4g}]",
                        )
                    else:
                        param_table.add_row(param_name)
                except Exception:
                    param_table.add_row(param_name, "[dim]error[/dim]", "", "")

            if has_more:
                param_table.add_row(
                    f"[dim]... and {n_params - max_params_display} more parameters[/dim]",
                    "",
                    "",
                    "" if show_ranges else None,
                )

            console.print(param_table)
            console.print()


def load_chains(
    filenames: List, labels: List | None = None, skip: None | float = 0.3
) -> dict:
    from .compute import Rminus1

    labels = [f"chain{i}" for i in range(len(filenames))] if labels is None else labels
    chains = {
        label: loadMCSamples(chain, settings={"ignore_rows": skip})
        for label, chain in zip(labels, filenames)
    }
    # After loading chains
    get_autocorr(chains)
    try:
        Rminus1(chains)
    except Exception:
        pass
    return chains


def get_chains(
    subset: float | str,
    config_file: str | Path = "chains.yaml",
    skip: float | None = None,
) -> dict:
    """
    Load MCMC chains based on kmin value from a YAML configuration file.

    Args:
        subset: Any key to determine which subset of chains to load
        config_file: Path to YAML configuration file. Default is "chains.yaml" in
            current directory. If not found, looks in the package root.
        skip: Fraction of samples to skip at beginning. If None, uses value from config.

    Returns:
        Dictionary of loaded chains with labels as keys and MCSamples objects as values

    Example:
        >>> chains = get_chains(1e-4)
        >>> chains = get_chains(LCDM, config_file="my_chains.yaml")
    """
    # Locate config file
    config_path = Path(config_file)
    if not config_path.exists():
        # Try package root directory
        package_root = Path(__file__).parent.parent.parent
        config_path = package_root / config_file
        if not config_path.exists():
            raise FileNotFoundError(
                f"Configuration file not found: {config_file}\n"
                f"Searched in current directory and {package_root}"
            )

    # Load YAML configuration
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)

    # Get root directory for chains
    root_dir = config.get("root", "")
    if root_dir:
        root_dir = Path(root_dir)
    else:
        root_dir = Path("")

    # Get default skip value if not provided
    if skip is None:
        skip = config.get("settings", {}).get("skip", 0.3)

    # Find matching kmin configuration
    chains_config = config.get("chains", {})
    subset_key = None

    # Try exact match first
    if subset in chains_config:
        subset_key = subset
    else:
        # Try string representation
        for key in chains_config.keys():
            if isinstance(key, str):
                try:
                    if float(key) == subset:
                        subset_key = key
                        break
                except (ValueError, TypeError):
                    pass
            elif key == subset:
                subset_key = key
                break

    if subset_key is None:
        available_subsets = list(chains_config.keys())
        raise ValueError(
            f"No configuration found for subset={subset}\n"
            f"Available subset values: {available_subsets}"
        )

    # Extract chain paths and labels
    chain_data = chains_config[subset_key]
    paths = []
    labels = []
    chains_with_h = []  # Track which chains need h parameter added

    for label, chain_info in chain_data.items():
        if isinstance(chain_info, dict):
            chain_path = root_dir / chain_info["path"]
            paths.append(str(chain_path))
            labels.append(label)
            if chain_info.get("add_h_parameter", False):
                chains_with_h.append(label)
        else:
            # Support simple format: label: path
            chain_path = root_dir / chain_info
            paths.append(str(chain_path))
            labels.append(label)

    tree = create_chain_tree(
        {label: chains_config.get(label, {}) for label in labels}, subset_name=subset
    )
    console.print(tree)

    # Load chains
    chains = load_chains(paths, labels, skip=skip)

    # Post-processing: Add h parameter for specific chains
    for label in chains_with_h:
        if label in chains:
            try:
                h_values = chains[label].getParams().H0 / 100
                chains[label].addDerived(h_values, "h", label="h")
                print_success(f"Added derived parameter 'h' to {label}")
            except Exception as e:
                print_warning(f"Could not add 'h' parameter to {label}", details=str(e))
                print(f"Warning: Could not add 'h' parameter to {label}: {e}")

    return chains
