from __future__ import annotations

import os
import json
import time
import math
from dataclasses import dataclass, asdict
from typing import Union, Sequence, Optional, Callable, Tuple, Dict, Any, List

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

_EPS = 1e-12


# =============================================================================
# Small utils
# =============================================================================

def _clip(x: float, lo: float, hi: float) -> float:
    return lo if x < lo else hi if x > hi else x


def _clip01(x: float) -> float:
    return _clip(float(x), 0.0, 1.0)


def _to_pred_list(preds: Optional[Sequence[Callable]]) -> list:
    if preds is None:
        return []
    if callable(preds):
        return [preds]
    try:
        return list(preds)
    except TypeError:
        return [preds]


def _normalize_components(components: Optional[Sequence[Union[int, str]]]) -> Optional[List[int]]:
    if components is None:
        return None
    comp_map = {"x": 0, "y": 1, "z": 2}
    out: List[int] = []
    for c in components:
        out.append(comp_map[c.lower()] if isinstance(c, str) else int(c))
    if any(c not in (0, 1, 2) for c in out):
        raise ValueError("freeze_components must be subset of {0,1,2,'x','y','z'}")
    return out


# =============================================================================
# Persistent state (global memory across separate runs)
# =============================================================================

@dataclass
class GalileoMCState:
    """
    Persistent MC adaptation state.

    Stored on disk to preserve tuned proposal scale across independent runs.

    Notes
    -----
    - step_A: current proposal std-dev scale (Å)
    - acc_ema: exponentially smoothed acceptance rate estimate
    - n_total: total number of MC proposals seen (for diagnostics)
    """
    step_A: float = 0.15
    acc_ema: float = 0.30
    n_total: int = 0
    updated_unix: float = 0.0


class _StateStore:
    """
    Lightweight JSON-backed state store.

    Keyed by user-provided `state_key` so you can maintain multiple
    independent adaptive memories (e.g., per system / per calculator).
    """
    def __init__(self, path: str, state_key: str = "default"):
        self.path = path
        self.state_key = state_key

    def load(self) -> GalileoMCState:
        try:
            with open(self.path, "r", encoding="utf-8") as f:
                blob = json.load(f)
            obj = blob.get(self.state_key, None)
            if obj is None:
                return GalileoMCState()
            return GalileoMCState(**obj)
        except FileNotFoundError:
            return GalileoMCState()
        except Exception:
            # Corrupted state file -> fail safe
            return GalileoMCState()

    def save(self, state: GalileoMCState) -> None:
        os.makedirs(os.path.dirname(self.path) or ".", exist_ok=True)
        blob: Dict[str, Any] = {}
        try:
            if os.path.exists(self.path):
                with open(self.path, "r", encoding="utf-8") as f:
                    blob = json.load(f)
            if not isinstance(blob, dict):
                blob = {}
        except Exception:
            blob = {}

        state.updated_unix = time.time()
        blob[self.state_key] = asdict(state)

        tmp = self.path + ".tmp"
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(blob, f, indent=2, sort_keys=True)
        os.replace(tmp, self.path)


# =============================================================================
# Constraint bridge (compatible with your predicate convention)
# =============================================================================

class _APMAdapter:
    def __init__(self, symbols, positions, cell, atomic_constraints=None):
        self.atomPositions = np.asarray(positions, float)
        self.atomLabelsList = np.asarray(symbols, dtype=object)
        _cell = np.asarray(cell, float) if cell is not None else None
        if _cell is None or _cell.shape != (3, 3) or not np.isfinite(_cell).all() or abs(np.linalg.det(_cell)) < 1e-9:
            _cell = np.eye(3, dtype=float)
        self.latticeVectors = _cell
        self.atomicConstraints = (np.asarray(atomic_constraints, bool)
                                  if atomic_constraints is not None else None)
        self.atomCount = len(self.atomPositions)


class _StructAdapter:
    def __init__(self, symbols, positions, cell, atomic_constraints=None):
        self.AtomPositionManager = _APMAdapter(symbols, positions, cell, atomic_constraints)


def _evaluate(idx: int, structure: _StructAdapter, constraints: Sequence[Callable], logic: str = "all") -> bool:
    if not constraints:
        return False
    if logic == "all":
        return all(c(idx, structure) for c in constraints)
    if logic == "any":
        return any(c(idx, structure) for c in constraints)
    raise ValueError("constraint_logic must be 'all' or 'any'")


def _select_indices_by_constraints(
    symbols,
    positions,
    cell,
    fixed: Optional[Sequence[bool]],
    constraints: Sequence[Callable],
    logic: str = "all",
) -> np.ndarray:
    N = len(symbols)
    adapter = _StructAdapter(symbols=symbols, positions=positions, cell=cell, atomic_constraints=fixed)
    sel = [i for i in range(N) if _evaluate(i, adapter, constraints, logic=logic)]
    return np.asarray(sel, dtype=int)


def _build_atoms(symbols: Sequence[str], positions: np.ndarray, cell: Optional[np.ndarray], calculator: object) -> Tuple[Atoms, bool]:
    if cell is not None:
        pbc = True
        cell = np.asarray(cell, float) if np.ndim(cell) == 2 else None
        if cell is None or cell.shape != (3, 3) or abs(np.linalg.det(cell)) < 1e-6:
            pbc, cell = False, None
    else:
        pbc = False

    atoms = Atoms(symbols=symbols, positions=positions, cell=cell, pbc=pbc)
    if calculator is not None:
        atoms.calc = calculator
    return atoms, pbc


def _movable_mask(N: int, selected: np.ndarray, action: str, comps: Optional[List[int]]) -> np.ndarray:
    """
    Boolean mask (N,3) where True means "allowed to move this DOF".
    """
    mask = np.ones((N, 3), dtype=bool)

    selected = np.asarray(selected, dtype=int)
    selected = np.unique(selected)

    if action not in ("freeze", "move_only"):
        raise ValueError("constraint_action must be 'freeze' or 'move_only'")

    if selected.size == 0:
        if action == "move_only":
            mask[:, :] = False
        return mask

    if action == "freeze":
        if comps is None:
            mask[selected, :] = False
        else:
            mask[selected[:, None], np.array(comps)[None, :]] = False
    else:
        mask[:, :] = False
        if comps is None:
            mask[selected, :] = True
        else:
            mask[selected[:, None], np.array(comps)[None, :]] = True

    return mask


# =============================================================================
# Main factory
# =============================================================================

def galileo_monte_carlo_fast(
    calculator: object = None,

    # --- core MC controls ---
    temperature_K: float = 300.0,
    n_steps: int = 500,

    # --- 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 safety ---
    Emax: Optional[float] = None,
    Emax_margin: float = 2.0,

    # --- 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 ---
    track_best: bool = True,
) -> 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.

    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()).

    Sampling correctness
    --------------------
    - Proposals are symmetric (Gaussian), so standard Metropolis acceptance applies.
    - If you require strict stationarity, set warmup_steps > 0 and enable_adaptation=True:
      adaptation occurs only during warmup, then step size is frozen.

    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 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")

    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) + _EPS)

    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,
        sampling_temperature: float = 0.0,   # API compatibility; not used unless you map external T -> K
        constraints: Optional[Sequence[Callable]] = None,
        output_filename: str = "mc_out.xyz",
        seed: Optional[int] = 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 movable_atoms.size == 0 or n_steps == 0:
            # Nothing to move (or no steps): return current evaluation
            E0 = float(atoms.get_potential_energy())
            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(E0),
                {
                    "mc_steps": int(n_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",
                },
            )

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

        # Initial energy, cap
        E = float(atoms.get_potential_energy())
        E_best = E
        pos_best = atoms.get_positions().copy()
        E_cap = float(Emax) if Emax is not None else float(E + Emax_margin)

        # ---------------------------------------------------------------------
        # Vectorized proposal generation (speed win)
        # ---------------------------------------------------------------------
        # Pre-generate:
        # - atom indices (into movable_atoms)
        # - Gaussian displacement vectors
        # - log-uniform random numbers for Metropolis test
        #
        # This reduces per-step Python overhead substantially.
        # ---------------------------------------------------------------------
        steps = int(n_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_Emax = 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} E={E:+.6f} eV acc={acc_rate:.3f} step={step_A:.3f}Å cap={E_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())

            # Optional hard barrier
            if E_new > E_cap:
                rejected_Emax += 1
                atoms.positions[a] = old
            else:
                dE = E_new - E

                # Log-Metropolis test
                # Accept if dE <= 0 OR log(u) < -beta*dE
                if dE <= 0.0 or (logu[i] < (-beta * dE)):
                    accepted += 1
                    accepted_window += 1
                    E = E_new

                    if track_best and E_new < E_best:
                        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 track_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_temperature_K": float(temperature_K),
            "mc_beta_eV": float(beta),
            "mc_steps": int(steps),
            "mc_acceptance": float(acc_rate),
            "mc_accepted": int(accepted),
            "mc_proposed": int(proposed),
            "mc_rejected_Emax": int(rejected_Emax),
            "mc_step_size_final_A": float(step_A),
            "mc_Emax_eV": float(E_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),
        }

        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