from typing import Any, Callable, Dict, List, Optional

import jax
import jax.numpy as jnp
import jax.flatten_util  # <-- Added this critical import
import numpy as np
from jaxtyping import PyTree

from inferix.nested import AbstractHostHypercubeNS
from inferix.result import Result

class PolyChord(AbstractHostHypercubeNS):
    """PolyChord Nested Sampling wrapper."""
    
    num_repeats: Optional[int] = None
    grade_dims: Optional[List[int]] = None
    grade_frac: Optional[List[float]] = None
    nlives: Optional[Dict[float, int]] = None
    cube_samples: Optional[Any] = None  
    
    nprior: int = -1
    nfail: int = -1
    do_clustering: bool = True
    feedback: int = 1
    precision_criterion: float = 0.001
    logzero: float = -1e30
    max_ndead: int = -1
    boost_posterior: float = 0.0
    posteriors: bool = True
    equals: bool = True
    cluster_posteriors: bool = True
    write_resume: bool = True
    write_paramnames: bool = False
    read_resume: bool = True
    write_stats: bool = True
    write_live: bool = True
    write_dead: bool = True
    write_prior: bool = True
    maximise: bool = False
    compression_factor: float = np.exp(-1.0)
    synchronous: bool = True
    base_dir: str = "chains"
    file_root: str = "test"
    seed: int = -1

    def run(
        self, 
        log_likelihood_fn: Callable, 
        prior_transform_fn: Callable, 
        y0: PyTree,  
        args: PyTree,
        *,
        nlive: int | None = None,
    ) -> Result:
        
        try:
            import pypolychord
            from pypolychord.settings import PolyChordSettings
        except ImportError:
            raise ImportError("pypolychord must be installed to use this sampler.")

        # 1. DERIVE GEOMETRY FROM y0
        flat_y0, reconstruct_fn = jax.flatten_util.ravel_pytree(y0)
        ndims = flat_y0.size

        _nlive = nlive if nlive is not None else ndims * 25
        _num_repeats = self.num_repeats if self.num_repeats is not None else ndims * 5
        _grade_dims = self.grade_dims if self.grade_dims is not None else [ndims]
        _grade_frac = self.grade_frac if self.grade_frac is not None else [1.0] * len(_grade_dims)
        _nlives = self.nlives if self.nlives is not None else {}

        settings = PolyChordSettings(nDims=ndims, nDerived=0)
        settings.nlive = _nlive
        settings.num_repeats = _num_repeats
        settings.nprior = self.nprior
        settings.nfail = self.nfail
        settings.do_clustering = self.do_clustering
        settings.feedback = self.feedback
        settings.precision_criterion = self.precision_criterion
        settings.logzero = self.logzero
        settings.max_ndead = self.max_ndead
        settings.boost_posterior = self.boost_posterior
        settings.posteriors = self.posteriors
        settings.equals = self.equals
        settings.cluster_posteriors = self.cluster_posteriors
        settings.write_resume = self.write_resume
        settings.write_paramnames = self.write_paramnames
        settings.read_resume = self.read_resume
        settings.write_stats = self.write_stats
        settings.write_live = self.write_live
        settings.write_dead = self.write_dead
        settings.write_prior = self.write_prior
        settings.maximise = self.maximise
        settings.compression_factor = self.compression_factor
        settings.synchronous = self.synchronous
        settings.base_dir = self.base_dir
        settings.file_root = self.file_root
        settings.seed = self.seed
        settings.grade_dims = _grade_dims
        settings.grade_frac = _grade_frac
        settings.nlives = _nlives
        settings.cube_samples = self.cube_samples

        # 2. JIT-COMPILED BRIDGES
        @jax.jit
        def jitted_likelihood(flat_theta_jax):
            struct_theta = reconstruct_fn(flat_theta_jax)
            return log_likelihood_fn(struct_theta, args)

        @jax.jit
        def jitted_prior(flat_u_jax):
            struct_u = reconstruct_fn(flat_u_jax)
            struct_theta = prior_transform_fn(struct_u, args)
            flat_theta, _ = jax.flatten_util.ravel_pytree(struct_theta)
            return flat_theta

        def polychord_likelihood(theta_np):
            logL = jitted_likelihood(jnp.asarray(theta_np))
            return float(logL), []

        def polychord_prior(u_np):
            return np.array(jitted_prior(jnp.asarray(u_np)))

        # 3. EXECUTE POLYCHORD
        output = pypolychord.run_polychord(
            loglikelihood=polychord_likelihood,
            nDims=ndims,
            nDerived=0,
            settings=settings,
            prior=polychord_prior
        )

        # 4. RESTRUCTURE RESULTS -> Returns a Batch of Caller PyTrees!
        structured_samples = jax.vmap(reconstruct_fn)(jnp.array(output.samples))

        return Result(
            samples=structured_samples, 
            final_state=None,
            logZ=jnp.array(output.logZ),
            logZ_err=jnp.array(output.logZerr),
            converged=jnp.array(True),
            result=RESULTS.successful,
            stats={"num_steps": len(output.samples)}
        )