from __future__ import annotations
import math
import copy
import os
import numpy as np
from typing import Any, List, Optional, Dict
from ..utils.chemistry import get_positions_species
from collections import Counter
from .base import TransitionKernel

class NestedSamplingKernel(TransitionKernel):
    """
    Transition kernel for Nested Sampling (NS).
    Extends the basic hard-threshold filter with statistical bookkeeping:
    X_k (prior volume), w_k (weights), and archive of dead points.
    """

    def __init__(self, cfg: Any = None, logger: Optional[Any] = None, ctx: Any = None):
        """
        Args:
            cfg: GA configuration object.
            logger: Optional logger for diagnostic messages.
            ctx: GA context object.
        """
        self.cfg = cfg
        self.logger = logger
        self.ctx = ctx
        
        # Consistent RNG source
        if ctx and hasattr(ctx, 'rng'):
            self.rng = ctx.rng
        else:
            self.rng = np.random.default_rng()
            if self.logger:
                self.logger.warning("[NestedSampling] No context RNG provided, using default_rng().")
        
        # NS Bookkeeping
        self.iteration = 0
        self.logX = 0.0  # log(X_0) = 0
        self.n_live = None
        self.step_interval = 1
        self.step_counter = 0
        self.internal_pop = []
        self.curr_phi = None # Cached energies of live points
        
        self.logZ = -np.inf
        self.H = 0.0
        self.rank_history = []
        self.compression_history = []
        self.accepted_replacements = 0
        self.rejected_candidates = 0
        self.failed_generations = 0
        self.tolerance = None
        self.track_diagnostics = True
        
        # Grand Canonical Support
        self.mu = {}
        if cfg and hasattr(cfg, 'ensemble') and cfg.ensemble.nested_sampling:
            ns_cfg = cfg.ensemble.nested_sampling
            self.mu = getattr(ns_cfg, 'chemical_potential', {})
            self.stochastic_shrinkage = getattr(ns_cfg, 'stochastic_shrinkage', False)
            self.strict_ns = getattr(ns_cfg, 'strict_ns', False)
            _beta = getattr(ns_cfg, 'beta', None)
            _T = getattr(ns_cfg, 'T', None)
            
            self.tolerance = getattr(ns_cfg, 'tolerance', None)
            self.track_diagnostics = getattr(ns_cfg, 'track_diagnostics', True)
            self.n_live = getattr(ns_cfg, 'n_live', None)
            self.step_interval = getattr(ns_cfg, 'step_interval', 1)
            
            if _beta is not None:
                self.beta = _beta
            elif _T is not None and _T > 0:
                # Physical inverse temperature 1/(kB*T)
                # Assumes energy in eV and T in Kelvin
                kB = 8.617333262145e-5
                self.beta = 1.0 / (kB * _T)
        
        # Determine output path
        self.output_path = "."
        if cfg and hasattr(cfg, 'output_path'):
            self.output_path = str(cfg.output_path)
        elif ctx and hasattr(ctx, 'output_path'):
            self.output_path = str(ctx.output_path)
            
        self.archive_path = os.path.join(self.output_path, "ns_dead_points.xyz")
        
        # Ensure directory exists
        if not os.path.exists(self.output_path):
            try:
                os.makedirs(self.output_path, exist_ok=True)
            except Exception:
                pass

    def calculate_log_weight(self, logX_prev: float, logX_curr: float) -> float:
        """w_k = X_{k-1} - X_k. Return as log weight."""
        # log_w = a + log1p(-exp(delta))
        delta = logX_curr - logX_prev
        
        # Guards for extreme deltas/mismatch
        if delta >= -1e-16:
            if self.logger:
                self.logger.warning(f"[NestedSampling] Warning: logX delta >= 0 ({delta}). Check compression rate.")
            return -np.inf
            
        # Clamp to avoid precision issues when delta is very small but non-zero
        delta = min(delta, -1e-16)
        
        if self.track_diagnostics:
            self.compression_history.append(delta)
        
        return logX_prev + math.log1p(-math.exp(delta))

    def step(
        self,
        parents: Any,
        candidates: List[Any],
        **kwargs
    ) -> List[Any]:
        """
        Executes Nested Sampling steps.
        Identifies successful candidates and replaces the worst points in the live set.
        """
        # Ensure we have lists of individuals
        if hasattr(parents, "containers"):
            parents_list = list(parents.containers)
        else:
            parents_list = list(parents) if not isinstance(parents, list) else parents
            
        candidates_list = list(candidates) if not isinstance(candidates, list) else candidates
        M = len(parents_list)

        if self.n_live is None:
            self.n_live = getattr(self.cfg.multiobjective, 'size', M) if self.cfg else M

        # State 1: Fill Internal Population initially
        if len(self.internal_pop) < self.n_live:
            for c in candidates_list:
                if c is not None and len(self.internal_pop) < self.n_live:
                    self.internal_pop.append(c)
            
            if len(self.internal_pop) == self.n_live:
                self.curr_phi = np.array([self._get_phi(p) for p in self.internal_pop], dtype=float)
                if self.logger:
                    self.logger.info(f"[NestedSampling] Internal population filled to {self.n_live}.")
                    
            # Return current candidates to keep GA stepping while filling
            ret_list = list(candidates_list)
            while len(ret_list) < M and len(ret_list) > 0:
                ret_list.append(self.rng.choice(ret_list))
            return ret_list[:M]

        # State 2: Active NS Evaluation
        self.step_counter += 1
        
        # 1. Get current worst energy threshold from internal alive population
        E_worst = np.max(self.curr_phi)
        
        # Export constraint explicitly to shared state for subsequent modular MCMC sampling
        if "shared_state" in kwargs:
            kwargs["shared_state"]["ns_phi_star"] = E_worst
        
        # 2. Collect successful candidates (those beating the internal worst)
        successes = []
        for child in candidates_list:
            if child is not None:
                E_child = self._get_phi(child)
                if E_child < E_worst:
                    successes.append((E_child, child))

        replaced = False
        
        # Execute NS Replacement only 1 out of every N times
        if self.step_counter % self.step_interval == 0:
            if successes:
                # Randomly pick ONE success to replace the worst internal point
                self.rng.shuffle(successes)
                E_child, child = successes[0]
                
                curr_worst_idx = np.argmax(self.curr_phi)
                E_dead = self.curr_phi[curr_worst_idx]
                
                # Archive the dead point
                self.iteration += 1
                logX_prev = self.logX
                
                if self.stochastic_shrinkage:
                    u = self.rng.random()
                    self.logX += math.log(max(u, 1e-300)) / self.n_live
                else:
                    self.logX -= 1.0 / self.n_live
                    
                delta = self.logX - logX_prev
                log_wk = self.calculate_log_weight(logX_prev, self.logX)
                
                rank = None
                if self.track_diagnostics:
                    self.compression_history.append(delta)
                    rank = np.sum(self.curr_phi < E_child)
                    self.rank_history.append(rank / self.n_live)
                
                # Evidence Accumulation
                logL = -self.beta * E_dead
                log_wL = log_wk + logL
                
                if self.track_diagnostics:
                    if np.isneginf(self.logZ):
                        self.logZ = log_wL
                        self.H = 0.0
                    else:
                        new_logZ = np.logaddexp(self.logZ, log_wL)
                        w1 = math.exp(self.logZ - new_logZ)
                        w2 = math.exp(log_wL - new_logZ)
                        self.H = w1 * (self.H + self.logZ) + w2 * logL - new_logZ
                        self.logZ = new_logZ
                
                # Replace in internal population cache
                self.curr_phi[curr_worst_idx] = E_child
                dead_point = self.internal_pop[curr_worst_idx]
                self._archive_dead_point(dead_point, E_dead, self.logX, log_wk, rank=rank, delta=delta, logZ=self.logZ, H=self.H)
                
                self.internal_pop[curr_worst_idx] = child
                
                replaced = True
                self.accepted_replacements += 1
            else:
                self.failed_generations += 1
                
            if self.logger:
                E_max_new = np.max(self.curr_phi)
                gen_str = kwargs.get('generation', '?')
                
                if replaced:
                    # Attempt to extract an ID from the structure
                    dead_id = getattr(dead_point, 'id', '?')
                    if dead_id == '?' and hasattr(dead_point, 'info'):
                        dead_id = dead_point.info.get('id', '?')
                    if dead_id == '?' and hasattr(dead_point, 'AtomPositionManager'):
                        meta = getattr(dead_point.AtomPositionManager, 'metadata', {})
                        if isinstance(meta, dict):
                            dead_id = meta.get('id', '?')
                            
                    status_str = f"Replaced (Eliminated ID: {dead_id}, E: {E_dead:.6f})"
                else:
                    status_str = "Rejected"
                    
                msg = f"[NestedSampling] Step {self.step_counter} (Gen {gen_str}): {status_str} | E_max={E_max_new:.6f} | logX={self.logX:.4f}"
                if self.track_diagnostics and not np.isneginf(self.logZ):
                    msg += f" | logZ={self.logZ:.4f} | H={self.H:.4f}"
                self.logger.info(msg)

                # NS Stopping Criterion check
                if self.tolerance is not None and not np.isneginf(self.logZ):
                    max_live_logL = -self.beta * np.min(self.curr_phi)
                    logZ_remaining = self.logX + max_live_logL
                    if logZ_remaining - self.logZ < -self.tolerance:
                        self.logger.info(f"[NestedSampling] Convergence met: logZ_remaining - logZ = {logZ_remaining - self.logZ:.4f} < -{self.tolerance}")
                        if self.ctx:
                            self.ctx.is_converge = True

        # Dynamic NS Hook Placeholder
        self._dynamic_ns_hook()

        # State 3: Construct Return Population
        # Return all successful candidates + complete with random drawings from internal_pop
        return_list = []
        for E_child, child in successes:
            if len(return_list) < M:
                return_list.append(child)
                
        if len(return_list) < M:
            remaining = M - len(return_list)
            idx_samples = self.rng.choice(len(self.internal_pop), size=remaining, replace=(remaining > len(self.internal_pop)))
            for idx in idx_samples:
                return_list.append(self.internal_pop[idx])
                
        return return_list

    def _dynamic_ns_hook(self):
        """
        Placeholder for Dynamic Nested Sampling (DNS) adjustments.
        Can evaluate posterior weights exp(log_wk + logL) and adjust 
        n_live dynamically in high-mass regions.
        """
        pass

    def _get_raw_energy(self, struct: Any) -> float:
        """Extracts the raw potential energy (E) from the structure."""
        # 1. Check AtomicData/APM metadata - prioritized in EZGA
        for obj in (struct, getattr(struct, "AtomPositionManager", None)):
            if obj is None: continue
            meta = getattr(obj, "metadata", {}) or {}
            if isinstance(meta, dict):
                for key in ("E_pot", "energy", "E"):
                    if key in meta:
                        return float(meta[key])
        
        # 2. Check direct attributes
        for obj in (struct, getattr(struct, "AtomPositionManager", None)):
            if obj is None: continue
            for attr in ("E", "energy", "E_pot"):
                val = getattr(obj, attr, None)
                if val is not None:
                    return float(val)
        
        # 3. Failure: Avoid silent 0.0 which corrupts NS ordering
        msg = f"Failed to extract energy from {type(struct)}"
        if self.logger:
            self.logger.error(msg)
        return float('inf')

    def _get_phi(self, struct: Any) -> float:
        """Calculates Grand Potential Phi = E - sum(mu_i * N_i)."""
        # 1. Prioritize pre-calculated Phi in metadata
        for obj in (struct, getattr(struct, "AtomPositionManager", None)):
            if obj is None: continue
            meta = getattr(obj, "metadata", {}) or {}
            if isinstance(meta, dict) and "Phi" in meta:
                return float(meta["Phi"])
        
        # 2. Otherwise calculate it from raw energy and species
        E = self._get_raw_energy(struct)
        if not self.mu:
            return E
            
        from ..utils.chemistry import get_positions_species
        try:
            apm = getattr(struct, "AtomPositionManager", struct)
            _, S = get_positions_species(apm)
        except Exception:
            return E
                
        phi = E
        counts = Counter(S)
        for species, mu_val in self.mu.items():
            n = counts.get(species, 0)
            phi -= mu_val * n
            
        return phi

    def _archive_dead_point(self, struct, phi_dead, logX, log_wk, rank=None, delta=None, logZ=None, H=None):
        """Appends the discarded structure to ns_dead_points.xyz."""
        try:
            E_raw = self._get_raw_energy(struct)
            apm = getattr(struct, "AtomPositionManager", struct)
            R, S = get_positions_species(apm)
            N = len(S)
            
            with open(self.archive_path, 'a') as f:
                f.write(f"{N}\n")
                
                # Ensure beta is not None
                beta_val = self.beta if self.beta is not None else 1.0
                logL = -beta_val * phi_dead
                
                # Metadata string for post-processing
                meta_parts = [
                    f"Iter={self.iteration}",
                    f"E_pot={E_raw:.10f}",
                    f"Phi={phi_dead:.10f}",
                    f"logL={logL:.12f}",
                    f"logX={logX:.12f}",
                    f"log_wk={log_wk:.12f}",
                    f"beta={beta_val:.6f}",
                    f"n_live={self.n_live}",
                    f"stochastic={int(self.stochastic_shrinkage)}"
                ]
                
                if rank is not None: meta_parts.append(f"rank={rank}")
                if delta is not None: meta_parts.append(f"delta={delta:.12f}")
                if logZ is not None and not np.isneginf(logZ): meta_parts.append(f"logZ={logZ:.12f}")
                if H is not None: meta_parts.append(f"H={H:.12f}")
                
                if self.mu:
                    mu_str = ",".join([f"{k}:{v}" for k, v in self.mu.items()])
                    meta_parts.append(f"mu={mu_str}")
                
                header = " ".join(meta_parts)
                f.write(header + "\n")

                for i in range(N):
                    f.write(f"{S[i]:<3} {R[i,0]:12.8f} {R[i,1]:12.8f} {R[i,2]:12.8f}\n")
        except Exception as e:
            if self.logger:
                self.logger.warning(f"[NestedSampling] Failed to archive dead point: {e}")
                # Log traceback for debugging if it's a TypeError or similar
                if isinstance(e, (TypeError, NameError, AttributeError)):
                    import traceback
                    self.logger.warning(traceback.format_exc())

    def allows_agent_sync(self) -> bool:
        """NS is tricky with sync, but generally fine if we treat immigrants as candidates."""
        return True
