from __future__ import annotations
import math
import copy
import os
import heapq
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.live_heap = []
        self._heap_counter = 0
        
        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
        
        self.stochastic_shrinkage = False
        self.strict_ns = False
        self.beta = None
        
        # 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)
            
            # Dynamic NS
            self.dynamic_ns = getattr(ns_cfg, 'dynamic_ns', False)
            self.dns_growth_factor = getattr(ns_cfg, 'dns_growth_factor', 1.5)
            self.dns_shrink_factor = getattr(ns_cfg, 'dns_shrink_factor', 0.5)
            self.dns_n_min = getattr(ns_cfg, 'dns_n_min', 10)
            self.dns_n_max = getattr(ns_cfg, 'dns_n_max', 10000)
            self.dns_peak_threshold = getattr(ns_cfg, 'dns_peak_threshold', 0.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)
            else:
                kB = 8.617333262145e-5
                self.beta = 1.0 / (kB * 300.0)

        # 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
                
        # Initialize default threshold in shared state for early simulations (State 1)
        if getattr(self, "ctx", None) is not None and hasattr(self.ctx, "shared_state"):
            self.ctx.shared_state.setdefault("ns_phi_star", float("inf"))
                
        # Truncate output file to prevent mixing incompatible runs
        open(self.archive_path, 'w').close()

    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)
        
        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.live_heap) < self.n_live:
            for c in candidates_list:
                if c is not None and len(self.live_heap) < self.n_live:
                    phi_c = self._get_phi(c)
                    self._heap_counter += 1
                    heapq.heappush(self.live_heap, (-phi_c, self._heap_counter, c))
            
            if len(self.live_heap) == self.n_live:
                if self.logger:
                    prefix = "[DNS]" if getattr(self, "dynamic_ns", False) else "[NestedSampling]"
                    self.logger.info(f"{prefix} Internal population filled to {self.n_live}.")
                    
            # Return current candidates to keep GA stepping while filling
            ret_list = list(candidates_list)
            if not ret_list:
                ret_list = [item[2] for item in self.live_heap]
                
            if ret_list:
                while len(ret_list) < M:
                    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
        phi_star = -self.live_heap[0][0]
        
        # Export constraint to global Context 'shared_state' for proper DI to simulators
        if getattr(self, "ctx", None) is not None and hasattr(self.ctx, "shared_state"):
            self.ctx.shared_state["ns_phi_star"] = float(phi_star)

        # 2. Collect successful candidates (those beating the internal worst)
        successes = []
        for child in candidates_list:
            if child is not None:
                phi_child = self._get_phi(child)
                if phi_child < phi_star:
                    successes.append((phi_child, child))
                else:
                    self.rejected_candidates += 1
        
        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)
                phi_child, child = successes[0]
                # Pop the worst point
                neg_phi_dead, _, dead_point = heapq.heappop(self.live_heap)
                phi_dead = -neg_phi_dead
                
                # 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 = sum(1 for item in self.live_heap if -item[0] < phi_child)
                    self.rank_history.append(rank / self.n_live)
                
                # Evidence Accumulation
                self._update_online_evidence(phi_dead, log_wk)
                
                self._archive_dead_point(dead_point, phi_dead, self.logX, log_wk, rank=rank, delta=delta, logZ=self.logZ, H=self.H)
                
                # Push the new child
                self._heap_counter += 1
                heapq.heappush(self.live_heap, (-phi_child, self._heap_counter, child))
                
                replaced = True
                self.accepted_replacements += 1
            else:
                self.failed_generations += 1
            
            if self.logger:
                phi_max_new = -self.live_heap[0][0]
                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}, Phi: {phi_dead:.6f})"
                else:
                    status_str = "Rejected"

                prefix = "[DNS]" if getattr(self, "dynamic_ns", False) else "[NestedSampling]"
                
                msg = f"{prefix} Step {self.step_counter} (Gen {gen_str}): {status_str} | Phi_max={phi_max_new:.6f} | n_live={self.n_live} | logX={self.logX:.4f}"
                
                if getattr(self, "dynamic_ns", False) and hasattr(self, "last_log_wL") and not np.isneginf(self.logZ):
                    W_k = math.exp(self.last_log_wL - self.logZ)
                    msg += f" | W_k={W_k:.2e}"
                    
                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
                self._check_thermo_convergence()

        # Dynamic NS Hook Placeholder
        self._dynamic_ns_hook()

        # State 3: Construct Return Population
        # Return all successful candidates + complete with random drawings from live_heap
        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.live_heap), size=remaining, replace=True)
            for idx in idx_samples:
                return_list.append(self.live_heap[idx][2])
                
        return return_list

    def _dynamic_ns_hook(self):
        """
        Dynamic Nested Sampling (DNS) adjustments.
        Evaluates posterior weights exp(log_wL - logZ) and adjusts n_live dynamically.
        """
        if not getattr(self, "dynamic_ns", False) or np.isneginf(self.logZ) or not hasattr(self, "last_log_wL"):
            return
            
        # Relative posterior fraction of the last discarded point
        W_k = math.exp(self.last_log_wL - self.logZ)
        
        # 1. Peak Thermodynamic Region - Increase walkers
        if W_k >= self.dns_peak_threshold and self.n_live < self.dns_n_max:
            new_n = min(self.dns_n_max, int(self.n_live * self.dns_growth_factor))
            if new_n > self.n_live:
                if self.logger:
                    self.logger.info(f"[DNS] Thermodynamic peak detected (W_k = {W_k:.3f}). Scaling UP: n_live {self.n_live} -> {new_n}")
                self.n_live = new_n

        # 2. Tail Thermodynamic Region - Decrease walkers
        elif W_k < (self.dns_peak_threshold * 1e-3) and self.n_live > self.dns_n_min:
            new_n = max(self.dns_n_min, int(self.n_live * self.dns_shrink_factor))
            if new_n < self.n_live:
                diff = self.n_live - new_n
                if self.logger:
                    self.logger.info(f"[DNS] Thermodynamic tail detected (W_k = {W_k:.2e}). Scaling DOWN: n_live {self.n_live} -> {new_n} (Dropping {diff} worst)")
                
                # Instantly drop the extra worst candidates
                self.n_live = new_n
                for _ in range(diff):
                    if len(self.live_heap) > self.n_live:
                        neg_phi_dead, _, dead_point = heapq.heappop(self.live_heap)
                        phi_dead = -neg_phi_dead
                        
                        self.iteration += 1
                        logX_prev = self.logX
                        
                        if self.stochastic_shrinkage:
                            u = self.rng.random()
                            self.logX += math.log(max(u, 1e-300)) / (len(self.live_heap) + 1)
                        else:
                            self.logX -= 1.0 / (len(self.live_heap) + 1)
                            
                        log_wk = self.calculate_log_weight(logX_prev, self.logX)
                        self._update_online_evidence(phi_dead, log_wk)
                        
                        rank = None
                        delta = self.logX - logX_prev
                        if self.track_diagnostics:
                            self.compression_history.append(delta)
                            self.rank_history.append(1.0)
                        
                        self._archive_dead_point(dead_point, phi_dead, self.logX, log_wk, rank=rank, delta=delta, logZ=self.logZ, H=self.H)
        
    def _update_online_evidence(self, phi_dead: float, log_wk: float):
        if self.beta is None:
            return

        logL = -self.beta * phi_dead
        log_wL = log_wk + logL
        self.last_log_wL = log_wL

        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
            
    def _check_thermo_convergence(self):
        if self.beta is None or self.tolerance is None or np.isneginf(self.logZ):
            return
            
        min_phi = min(-item[0] for item in self.live_heap) if self.live_heap else float('inf')
        max_live_logL = -self.beta * min_phi
        logZ_remaining = self.logX + max_live_logL
        if logZ_remaining - self.logZ < -self.tolerance:
            if self.logger:
                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

    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.strict_ns:
            raise ValueError(msg)
        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")
                
                # Metadata string for post-processing
                meta_parts = [
                    f"Iter={self.iteration}",
                    f"E_pot={E_raw:.10f}",
                    f"Phi={phi_dead:.10f}",
                    f"logX={logX:.12f}",
                    f"log_wk={log_wk:.12f}",
                    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 self.mu:
                    mu_str = ",".join([f"{k}:{v}" for k, v in self.mu.items()])
                    meta_parts.append(f"mu={mu_str}")
                    
                _beta = getattr(self, 'beta', None)
                if _beta is not None:
                    logL = -_beta * phi_dead
                    meta_parts.append(f"logL={logL:.12f}")
                    meta_parts.append(f"beta={_beta:.6f}")
                    
                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}")
                
                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

def reweight_ns_dead_points(candidates: Any, T: float) -> Dict[str, float]:
    """
    Post-processes a collection of Nested Sampling dead points to calculate 
    thermodynamic observables at a new temperature T.
    
    Args:
        candidates: Iterable or container of structures (dead points) with NS metadata.
        T: New temperature in Kelvin.
        
    Returns:
        Dictionary containing T, beta, logZ, E_mean, Phi_mean, H, and Omega.
    """
    if hasattr(candidates, "containers"):
        candidates_list = list(candidates.containers)
    else:
        candidates_list = list(candidates) if not isinstance(candidates, list) else candidates
        
    phi_list = []
    log_wk_list = []
    E_pot_list = []
    
    # Step 3: Define beta from T
    kB = 8.617333262145e-5  # eV/K
    beta = 1.0 / (kB * T)
    
    # Step 2: Extract from each structure
    for struct in candidates_list:
        meta = {}
        if hasattr(struct, "info"):
            meta.update(struct.info)
            
        obj = getattr(struct, "AtomPositionManager", struct)
        if hasattr(obj, "metadata") and isinstance(obj.metadata, dict):
            meta.update(obj.metadata)
            
        try:
            phi_val = float(meta["Phi"])
            log_wk_val = float(meta["log_wk"])
            
            if "E_pot" in meta:
                E_val = float(meta["E_pot"])
            elif "energy" in meta:
                E_val = float(meta["energy"])
            elif "E" in meta:
                E_val = float(meta["E"])
            else:
                E_val = phi_val

            phi_list.append(phi_val)
            log_wk_list.append(log_wk_val)
            E_pot_list.append(E_val)
        except (KeyError, ValueError, TypeError):
            continue
            
    if not phi_list:
        raise ValueError("No valid dead points with 'Phi' and 'log_wk' metadata found.")
        
    phi_array = np.array(phi_list, dtype=np.float64)
    log_wk_array = np.array(log_wk_list, dtype=np.float64)
    E_array = np.array(E_pot_list, dtype=np.float64)
    
    # Step 4: Calculate thermodynamic log-weights
    log_terms = log_wk_array - beta * phi_array
    
    # Step 5: Stable logZ using log-sum-exp
    logZ = np.logaddexp.reduce(log_terms)
    
    # Step 6: Normalized posterior weights
    log_post = log_terms - logZ
    post = np.exp(log_post)
    
    # Step 7: Calculate observables
    E_mean = np.sum(post * E_array)
    Phi_mean = np.sum(post * phi_array)
    
    valid_mask = post > 0
    H = -np.sum(post[valid_mask] * log_post[valid_mask])
    
    Omega = -logZ / beta
    
    return {
        "T": float(T),
        "beta": float(beta),
        "logZ": float(logZ),
        "E_mean": float(E_mean),
        "Phi_mean": float(Phi_mean),
        "H": float(H),
        "Omega": float(Omega)
    }

