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 galilean_constrained_mc(
    calculator: object = None,

    # --- GMC controls ---
    constraint_kind: str = "energy",    # "energy", "phi_fixed_composition", or "custom"
    observable_fn: Optional[Callable] = None,
    observable_grad_fn: Optional[Callable] = 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,
    
    eps_A: float = 0.05,
    L_traj: int = 8,
    N_traj: int = 10,

    # --- persistent adaptive proposal memory ---
    state_path: str = "gmc_state.json",
    state_key: str = "default",
    enable_adaptation: bool = False,
    warmup_trajectories: int = 0,

    # --- adaptation control (Robbins–Monro style for eps_A)
    target_acceptance: float = 0.30,
    adapt_interval: int = 5,
    adapt_gain: float = 0.35,
    eps_bounds_A: Tuple[float, float] = (0.01, 0.80),
    acc_ema_alpha: float = 0.05,

    # --- 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,
) -> Callable[..., Tuple[np.ndarray, np.ndarray, Optional[np.ndarray], float, Dict[str, Any]]]:
    """
    True Galilean Monte Carlo sampling for exact Nested Sampling replacements.
    
    This uses:
    - active subspace velocities and linear propagation.
    - specular reflections at the constraint boundary based on the gradient.
    - multiple sub-steps to form constrained trajectories.
    """
    
    #if observable_max is None:
    #    raise ValueError("observable_max is required for Galilean MC (can be float or Callable)")
    if constraint_kind not in ("energy", "phi_fixed_composition", "custom"):
        raise ValueError(f"Unknown constraint_kind: {constraint_kind}")
    if constraint_kind == "custom" and (observable_fn is None or observable_grad_fn is None):
        raise ValueError("custom constraint requires both observable_fn and observable_grad_fn")

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

    store = _StateStore(path=state_path, state_key=state_key)
    state = store.load()
    if state.n_total == 0:
        state.step_A = float(eps_A)
        state.acc_ema = float(target_acceptance)

    eps_lo, eps_hi = eps_bounds_A
    state.step_A = float(_clip(state.step_A, eps_lo, eps_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 = "gmc_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)
        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)

        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)

        # Build active subspace
        active_indices = np.flatnonzero(movable.reshape(-1))
        n_active = active_indices.size

        # ---------------------------------------------------------------------
        # Dynamic Constraint Resolution (Observer pattern via Dependency Injection)
        # ---------------------------------------------------------------------
        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 obs_cap is None:
            raise ValueError("Could not resolve observable_max cap dynamically in GMC. Check shared_state or observable_max callable.")

        if n_active == 0 or N_traj == 0 or L_traj == 0:
            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),
                {"gmc_note": "no movable atoms or zero trajectories", "gmc_eps": float(state.step_A)},
            )

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

        def get_val_and_grad(at: Any) -> Tuple[float, np.ndarray]:
            if constraint_kind in ("energy", "phi_fixed_composition"):
                v = float(at.get_potential_energy())
                f = np.asarray(at.get_forces(), float).reshape(-1)
                g = -f  # gradient of energy
            else:
                v = float(observable_fn(at))
                g = np.asarray(observable_grad_fn(at), float).reshape(-1)
            
            return v, g[active_indices]

        E_initial = float(atoms.get_potential_energy())
        x_start_flat = atoms.get_positions().copy().reshape(-1)
        x_active = x_start_flat[active_indices].copy()

        # Update atoms view
        def set_active_pos(x_vec: np.ndarray):
            x_flat = atoms.get_positions().reshape(-1)
            x_flat[active_indices] = x_vec
            atoms.set_positions(x_flat.reshape(-1, 3))

        eps_current = float(state.step_A)
        
        # Stats
        bounces = 0
        accepted_substeps = 0
        proposed_substeps = 0
        rejected_traj = 0  # number of full trajectories abandoned
        
        accepted_window = 0
        proposed_window = 0

        t0 = time.time()

        adapt_until = N_traj if (warmup_trajectories <= 0) else min(int(warmup_trajectories), N_traj)

        for t in range(N_traj):
            # Sample uniform random velocity in active subspace
            v = rng.normal(size=n_active)
            v /= max(np.linalg.norm(v), 1e-15)
            
            x_traj_start = x_active.copy()

            for l in range(L_traj):
                proposed_substeps += 1
                proposed_window += 1
                
                x_trial = x_active + eps_current * v
                set_active_pos(x_trial)
                obs_trial, grad_active = get_val_and_grad(atoms)
                
                if obs_trial < obs_cap:
                    # Valid step
                    x_active = x_trial
                    accepted_substeps += 1
                    accepted_window += 1
                else:
                    # Reflection
                    norm = np.linalg.norm(grad_active)
                    if norm < 1e-12:
                        # Degenerate gradient, cannot reflect gracefully
                        set_active_pos(x_active)
                        v = -v
                    else:
                        n = grad_active / norm
                        v_ref = v - 2.0 * np.dot(v, n) * n
                        
                        x_ref_trial = x_active + eps_current * v_ref
                        set_active_pos(x_ref_trial)
                        obs_ref, _ = get_val_and_grad(atoms)
                        
                        if obs_ref < obs_cap:
                            x_active = x_ref_trial
                            v = v_ref
                            accepted_substeps += 1
                            accepted_window += 1
                            bounces += 1
                        else:
                            # Both trial and reflected point fell outside
                            # Reverse velocity as a fallback
                            set_active_pos(x_active)
                            v = -v

                if write_interval and (l + t * L_traj) % write_interval == 0:
                    try: ase.io.write(output_filename, atoms, format="extxyz", append=True)
                    except Exception: pass

            if log_interval and (t % log_interval == 0):
                acc_r = accepted_substeps / max(proposed_substeps, 1)
                print(f"[GMC] t={t:4d} acc={acc_r:.3f} bounces={bounces} eps={eps_current:.4f}Å", flush=True)

            # Adapt EPS at ends of trajectories during warmup
            if enable_adaptation and (t + 1) % adapt_interval == 0 and (t + 1) <= adapt_until:
                acc_w = accepted_window / max(proposed_window, 1)
                eps_current *= math.exp(adapt_gain * (acc_w - target_acceptance))
                eps_current = float(_clip(eps_current, eps_lo, eps_hi))
                accepted_window = 0
                proposed_window = 0

        # Always return the final state, no track_best.
        set_active_pos(x_active)
        E_out = float(atoms.get_potential_energy())
        
        try: ase.io.write(output_filename, atoms, format="extxyz", append=True)
        except Exception: pass

        acc_rate = accepted_substeps / max(proposed_substeps, 1)
        state.step_A = float(eps_current)
        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_substeps)
        store.save(state)

        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] = {
            "gmc_eps_final_A": float(eps_current),
            "gmc_trajectories": int(N_traj),
            "gmc_steps_per_traj": int(L_traj),
            "gmc_proposed_substeps": int(proposed_substeps),
            "gmc_accepted_substeps": int(accepted_substeps),
            "gmc_acceptance": float(acc_rate),
            "gmc_bounces": int(bounces),
            "gmc_cap": float(obs_cap),
            "gmc_state_acc_ema": float(state.acc_ema),
            "gmc_runtime_s": float(time.time() - t0),
            "gmc_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