import os
import time
import math
from typing import Union, Sequence, Optional, Callable, Tuple, Dict, Any, List

import numpy as np
from ase import Atoms
import ase.io

from ezga.simulator.galileo_mc import _EPS, _clip, _normalize_components, _to_pred_list, _StateStore, _build_atoms, _select_indices_by_constraints, _movable_mask

def adaptive_constrained_metropolis_mc(
    calculator: object = None,

    # --- core MC controls ---
    mode: str = "canonical_truncated",   # "canonical_truncated" or "nested_sampling"
    temperature_K: float = 300.0,
    n_steps: int = 500,
    n_sweeps: Optional[float] = None,    # If provided, steps = ceil(n_sweeps * movable_atoms)

    # --- persistent adaptive proposal memory ---
    state_path: str = "mc_state.json",
    state_key: str = "default",           # e.g. system id / hash / calculator name
    enable_adaptation: bool = True,
    warmup_steps: int = 0,               # if >0: adapt only in first warmup_steps for stricter sampling after

    # --- proposal magnitude ---
    max_displacement_A: float = 0.15,    # used only if no state exists
    displacement_bounds_A: Tuple[float, float] = (0.01, 0.80),

    # --- adaptation control (Robbins–Monro style) ---
    target_acceptance: float = 0.30,
    adapt_interval: int = 25,
    adapt_gain: float = 0.35,
    acc_ema_alpha: float = 0.05,         # persistent acceptance EMA update

    # --- energy/observable safety ---
    Emax: Optional[float] = None,
    Emax_margin: float = 2.0,            # used only in canonical_truncated if Emax is None
    observable_fn: Optional[Callable[[Atoms], float]] = None,
    # observable_max can be a static float, or a Callable that returns a float. 
    # The actual constraint is resolved at runtime per-run, preferring shared_state['ns_phi_star'] if present.
    observable_max: Optional[Union[float, Callable[[], float]]] = None, 

    # --- IO knobs ---
    log_interval: int = 0,
    write_interval: int = 0,

    # --- constraints ---
    constraint_logic: str = "all",
    constraint_action: str = "freeze",
    freeze_components: Optional[Sequence[Union[int, str]]] = None,
    constraints: Optional[Sequence[Callable]] = None,

    # --- behavior ---
    return_best: bool = False,
) -> Callable[..., Tuple[np.ndarray, np.ndarray, Optional[np.ndarray], float, Dict[str, Any]]]:
    """
    High-throughput, constraint-aware Metropolis MC (ASE backend) with persistent adaptive step size.

    Modes:
    - canonical_truncated: Standard Metropolis MC, targeting Boltzmann-weighted distribution below an energy cap.
    - nested_sampling: Unbiased hard-threshold MC sampling, exact for NS uniform exploration of the constraint prior.

    Key performance choices
    -----------------------
    - Single-atom moves (high acceptance per energy evaluation).
    - Vectorized proposal generation: precompute atom indices, displacements, and log(u).
      This reduces Python overhead and RNG calls inside the loop.
    - Log-Metropolis test: log(u) < -βΔE (avoids exp()).

    Persistent memory
    -----------------
    - Step size and acceptance EMA are stored in JSON under `state_key`.
    - This makes the proposal scale immediately reasonable in future runs.
    """

    if mode not in ("canonical_truncated", "nested_sampling"):
        raise ValueError("mode must be 'canonical_truncated' or 'nested_sampling'")
    if temperature_K <= 0:
        raise ValueError("temperature_K must be > 0")
    if n_steps < 0:
        raise ValueError("n_steps must be >= 0")
    if adapt_interval <= 0:
        raise ValueError("adapt_interval must be > 0")
    if mode == "nested_sampling" and observable_max is None:
        raise ValueError("observable_max is required for nested_sampling mode (can be float or Callable)")

    freeze_components_norm = _normalize_components(freeze_components)
    factory_constraints = _to_pred_list(constraints)

    # Metropolis beta in eV^-1
    kB_eV_per_K = 8.617333262145e-5
    beta = 1.0 / (kB_eV_per_K * float(temperature_K))
    if mode == "nested_sampling":
        beta = 0.0

    store = _StateStore(path=state_path, state_key=state_key)
    state = store.load()
    if state.n_total == 0:
        # First time: initialize from config
        state.step_A = float(max_displacement_A)
        state.acc_ema = float(target_acceptance)

    step_lo, step_hi = displacement_bounds_A
    state.step_A = float(_clip(state.step_A, step_lo, step_hi))

    def run(
        symbols: Union[np.ndarray, Sequence[str]],
        positions: np.ndarray,
        cell: Optional[np.ndarray],
        *,
        fixed: Optional[np.ndarray] = None,
        constraints: Optional[Sequence[Callable]] = None,
        output_filename: str = "mc_out.xyz",
        seed: Optional[int] = None,
        shared_state: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray], float, Dict[str, Any]]:

        rng = np.random.default_rng(seed)

        symbols_arr = np.asarray(symbols, dtype=object)
        positions_arr = np.asarray(positions, dtype=float)
        if positions_arr.ndim != 2 or positions_arr.shape[1] != 3:
            raise ValueError("positions must be (N,3)")
        cell_arr = np.asarray(cell, dtype=float) if cell is not None else None

        atoms, pbc_flag = _build_atoms(symbols_arr, positions_arr, cell_arr, calculator)

        # Merge constraint predicates
        combined_constraints = list(factory_constraints)
        combined_constraints.extend(_to_pred_list(constraints))

        adapter_cell = cell_arr if (pbc_flag and cell_arr is not None) else np.eye(3, dtype=float)
        if combined_constraints:
            selected = _select_indices_by_constraints(
                symbols_arr, positions_arr, adapter_cell, fixed,
                combined_constraints, logic=constraint_logic
            )
        else:
            selected = np.array([], dtype=int)

        N = len(atoms)
        movable = _movable_mask(N, selected, constraint_action, freeze_components_norm)

        # Select movable atoms (at least one movable DOF)
        movable_atoms = np.where(movable.any(axis=1))[0]
        
        if n_sweeps is not None:
            actual_steps = int(np.ceil(n_sweeps * movable_atoms.size))
        else:
            actual_steps = int(n_steps)

        # ---------------------------------------------------------------------
        # Dynamic Constraint Resolution (Observer pattern via Dependency Injection)
        # ---------------------------------------------------------------------
        # Order of precedence for resolving the threshold cap:
        # 1. `shared_state["ns_phi_star"]` injected by High-Level algorithms (e.g., NestedSamplingKernel)
        # 2. `observable_max()` evaluated as a Callable right now
        # 3. `observable_max` evaluated as a static float
        # 4. `Emax` (for canonical mode defaults)
        obs_cap = None
        if shared_state is not None and "ns_phi_star" in shared_state:
            obs_cap = float(shared_state["ns_phi_star"])
        elif observable_max is not None:
            obs_cap = float(observable_max()) if callable(observable_max) else float(observable_max)
        
        if mode == "nested_sampling" and obs_cap is None:
            raise ValueError("Could not resolve observable_max for nested_sampling. Check shared_state or observable_max callable.")
        
        # Initial energy
        E = float(atoms.get_potential_energy())
        if observable_fn is not None:
            obs = float(observable_fn(atoms))
            obs_name = "custom"
        else:
            obs = E
            obs_name = "E"

        if mode != "nested_sampling" and obs_cap is None:
            if Emax is not None:
                obs_cap = float(Emax)
            else:
                obs_cap = float(E + Emax_margin)
                
        if movable_atoms.size == 0 or actual_steps == 0:
            # Nothing to move (or no steps): return current evaluation
            cell_obj = atoms.get_cell()
            pbc_out = bool(atoms.get_pbc().any())
            cell_out = cell_obj.array if (pbc_out and float(cell_obj.volume) > 1e-12) else None
            return (
                np.asarray(atoms.get_positions(), float),
                np.asarray(atoms.get_chemical_symbols(), object),
                np.asarray(cell_out, float) if cell_out is not None else None,
                float(E),
                {
                    "mc_steps": int(actual_steps),
                    "mc_acceptance": 0.0,
                    "mc_step_size_final_A": float(state.step_A),
                    "mc_beta_eV": float(beta),
                    "mc_note": "no movable atoms or zero steps",
                    "mc_mode": str(mode),
                    "mc_cap": float(obs_cap) if obs_cap is not None else -1.0,
                    "mc_return_state": "best" if return_best else "final",
                },
            )

        # Output dir
        out_dir = os.path.dirname(output_filename) or "."
        os.makedirs(out_dir, exist_ok=True)

        E_best = E
        obs_best = obs
        pos_best = atoms.get_positions().copy()

        # ---------------------------------------------------------------------
        # Vectorized proposal generation (speed win)
        # ---------------------------------------------------------------------
        steps = int(actual_steps)
        pick = rng.integers(0, movable_atoms.size, size=steps, endpoint=False)
        atom_idx = movable_atoms[pick]  # shape (steps,)
        disp = rng.normal(0.0, 1.0, size=(steps, 3))  # unit gaussian
        logu = np.log(rng.random(size=steps) + _EPS)  # log(u) in (-inf, 0]

        # Stats
        accepted = 0
        proposed = 0
        rejected_boundary = 0
        rejected_metropolis = 0
        accepted_window = 0
        proposed_window = 0

        step_A = float(state.step_A)
        t0 = time.time()

        def _maybe_log(i: int):
            if log_interval and (i % log_interval == 0):
                acc_rate = accepted / max(proposed, 1)
                print(
                    f"[MC] i={i:6d} {obs_name}={obs:+.6f} acc={acc_rate:.3f} step={step_A:.3f}Å cap={obs_cap:+.3f}",
                    flush=True,
                )

        def _maybe_write(i: int):
            if write_interval and (i % write_interval == 0):
                try:
                    ase.io.write(output_filename, atoms, format="extxyz", append=True)
                except Exception:
                    pass

        # Decide adaptation horizon (warmup)
        adapt_until = steps if (warmup_steps <= 0) else min(int(warmup_steps), steps)

        # MC loop (single expensive call: get_potential_energy per step)
        for i in range(steps):
            proposed += 1
            proposed_window += 1

            a = int(atom_idx[i])
            old = atoms.positions[a].copy()

            # Apply displacement only on allowed components for this atom
            d = disp[i] * step_A
            d *= movable[a]  # zero frozen components

            atoms.positions[a] = old + d

            E_new = float(atoms.get_potential_energy())
            if observable_fn is not None:
                obs_new = float(observable_fn(atoms))
            else:
                obs_new = E_new

            # Hard barrier rejection
            if obs_new >= obs_cap:
                rejected_boundary += 1
                atoms.positions[a] = old
            else:
                accepted_move = False
                if mode == "nested_sampling":
                    # Hard-threshold already passed
                    accepted_move = True
                else:
                    dE = E_new - E
                    # Log-Metropolis test
                    if dE <= 0.0 or (logu[i] < (-beta * dE)):
                        accepted_move = True
                    else:
                        rejected_metropolis += 1

                if accepted_move:
                    accepted += 1
                    accepted_window += 1
                    E = E_new
                    obs = obs_new

                    if return_best and obs_new < obs_best:
                        obs_best = obs_new
                        E_best = E_new
                        pos_best = atoms.get_positions().copy()
                else:
                    atoms.positions[a] = old

            # Adapt step size (only during warmup if warmup_steps>0)
            if enable_adaptation and (i + 1) % adapt_interval == 0 and (i + 1) <= adapt_until:
                acc_w = accepted_window / max(proposed_window, 1)

                # Multiplicative Robbins–Monro update (scale-free)
                # step <- step * exp(gain*(acc-target))
                step_A *= math.exp(adapt_gain * (acc_w - target_acceptance))
                step_A = float(_clip(step_A, step_lo, step_hi))

                accepted_window = 0
                proposed_window = 0

            _maybe_log(i)
            _maybe_write(i)

        # Export final/best (best-effort)
        if return_best:
            atoms.set_positions(pos_best)
            E_out = float(E_best)
        else:
            E_out = float(E)

        try:
            ase.io.write(output_filename, atoms, format="extxyz", append=True)
        except Exception:
            pass

        # Persist global memory
        acc_rate = accepted / max(proposed, 1)
        state.step_A = float(step_A)
        state.acc_ema = float((1.0 - acc_ema_alpha) * state.acc_ema + acc_ema_alpha * acc_rate)
        state.n_total = int(state.n_total + proposed)
        store.save(state)

        # Cell output
        cell_obj = atoms.get_cell()
        pbc_out = bool(atoms.get_pbc().any())
        cell_out = cell_obj.array if (pbc_out and float(cell_obj.volume) > 1e-12) else None

        meta: Dict[str, Any] = {
            "mc_mode": str(mode),
            "mc_return_state": "best" if return_best else "final",
            "mc_constraint_quantity": str(obs_name),
            "mc_temperature_K": float(temperature_K) if mode != "nested_sampling" else 0.0,
            "mc_beta_eV": float(beta),
            "mc_steps": int(steps),
            "mc_sweeps_requested": float(n_sweeps) if n_sweeps is not None else 0.0,
            "mc_acceptance": float(acc_rate),
            "mc_accepted": int(accepted),
            "mc_proposed": int(proposed),
            "mc_rejected_boundary": int(rejected_boundary),
            "mc_rejected_metropolis": int(rejected_metropolis),
            "mc_step_size_final_A": float(step_A),
            "mc_cap": float(obs_cap),
            "mc_state_path": str(state_path),
            "mc_state_key": str(state_key),
            "mc_state_acc_ema": float(state.acc_ema),
            "mc_state_total_proposals": int(state.n_total),
            "mc_runtime_s": float(time.time() - t0),
            "mc_adaptation_enabled": bool(enable_adaptation),
            "mc_warmup_steps": int(warmup_steps),
            "mc_seed": int(seed) if seed is not None else -1,
        }

        return (
            np.asarray(atoms.get_positions(), float),
            np.asarray(atoms.get_chemical_symbols(), object),
            np.asarray(cell_out, float) if cell_out is not None else None,
            float(E_out),
            meta,
        )

    return run
