from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable, Optional, Sequence

import numpy as np

K_BOLTZ = 0.008314462618  # kJ/mol/K


@dataclass
class GistConfig:
    grid_spacing: float = 0.1
    padding: float = 0.5
    temperature: float = 300.0
    length_scale: float = 0.1
    cutoff: float = 1.0
    orientation_bins: int = 12
    chunk_frames: Optional[int] = None
    frame_indices: Optional[Sequence[int]] = None
    solute_selection: Optional[str] = None
    max_frames: Optional[int] = None
    water_resnames: Sequence[str] = ("HOH", "WAT", "SOL", "TIP3", "OPC")
    bulk_density: Optional[float] = None
    energy_method: str = "direct"  # "direct" or "pme"


@dataclass
class GistResult:
    density: np.ndarray
    energy_sw: np.ndarray
    energy_ww: np.ndarray
    ts_trans: np.ndarray
    ts_orient: np.ndarray
    counts: np.ndarray
    origin: np.ndarray
    spacing: float
    n_frames: int
    voxel_volume: float
    bulk_density: float

    def as_dict(self):
        return {
            "density": self.density,
            "energy_sw": self.energy_sw,
            "energy_ww": self.energy_ww,
            "ts_trans": self.ts_trans,
            "ts_orient": self.ts_orient,
            "counts": self.counts,
            "origin": self.origin,
            "spacing": self.spacing,
            "n_frames": self.n_frames,
            "voxel_volume": self.voxel_volume,
            "bulk_density": self.bulk_density,
        }


def _require_openmm():
    try:
        from openmm import unit, openmm as mm
    except Exception as exc:  # pragma: no cover
        raise ImportError("openmm is required for gist") from exc
    return mm, unit


def _find_nonbonded(system):
    mm, _ = _require_openmm()
    for force in system.getForces():
        if isinstance(force, mm.NonbondedForce):
            return force
    raise ValueError("openmm_system has no NonbondedForce")


def _build_custom_system(openmm_system, cutoff):
    mm, unit = _require_openmm()
    nb = _find_nonbonded(openmm_system)
    system = mm.System()
    for i in range(openmm_system.getNumParticles()):
        system.addParticle(openmm_system.getParticleMass(i))

    energy_expr = (
        "4*epsilon*((sigma/r)^12 - (sigma/r)^6) + ONE_4PI_EPS0*charge1*charge2/r"
    )
    force = mm.CustomNonbondedForce(energy_expr)
    force.addPerParticleParameter("charge")
    force.addPerParticleParameter("sigma")
    force.addPerParticleParameter("epsilon")
    if openmm_system.usesPeriodicBoundaryConditions():
        force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
    else:
        force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffNonPeriodic)
    force.setCutoffDistance(cutoff * unit.nanometer)
    force.setUseLongRangeCorrection(False)

    for i in range(nb.getNumParticles()):
        charge, sigma, epsilon = nb.getParticleParameters(i)
        force.addParticle([charge, sigma, epsilon])

    for i in range(nb.getNumExceptions()):
        p1, p2, _q, _sig, _eps = nb.getExceptionParameters(i)
        force.addExclusion(p1, p2)

    system.addForce(force)
    return system, force


def _water_groups(topology, water_resnames: Iterable[str]):
    resnames = {r.upper() for r in water_resnames}
    waters = []
    atom_indices = {atom: idx for idx, atom in enumerate(topology.atoms())}
    for res in topology.residues():
        if res.name.upper() not in resnames:
            continue
        atoms = list(res.atoms())
        indices = [atom_indices[a] for a in atoms]
        oxy = None
        hydrogens = []
        for atom in atoms:
            name = atom.name.upper()
            if atom.element is not None and atom.element.symbol.upper() == "O":
                oxy = atom_indices[atom]
            elif name.startswith("O"):
                oxy = atom_indices[atom]
            elif atom.element is not None and atom.element.symbol.upper() == "H":
                hydrogens.append(atom_indices[atom])
            elif name.startswith("H"):
                hydrogens.append(atom_indices[atom])
        waters.append((indices, oxy, hydrogens))
    return waters


def _grid_from_positions(positions, spacing, padding):
    min_xyz = positions.min(axis=0) - padding
    max_xyz = positions.max(axis=0) + padding
    dims = np.maximum(np.ceil((max_xyz - min_xyz) / spacing).astype(int), 1)
    return min_xyz, dims


def _voxel_index(pos, origin, spacing, dims):
    idx = np.floor((pos - origin) / spacing).astype(int)
    if np.any(idx < 0) or np.any(idx >= dims):
        return None
    return tuple(idx.tolist())


class _EnergyPartitioner:
    def __init__(self, openmm_system, cutoff, platform_name=None):
        mm, unit = _require_openmm()
        self.system, self.force = _build_custom_system(openmm_system, cutoff)
        self.integrator = mm.VerletIntegrator(1.0 * unit.femtosecond)
        if platform_name:
            platform = mm.Platform.getPlatformByName(platform_name)
            self.context = mm.Context(self.system, self.integrator, platform)
        else:
            self.context = mm.Context(self.system, self.integrator)
        self.unit = unit

    def set_positions(self, positions_nm, box=None):
        if box is not None:
            lx, ly, lz = box
            a = self.unit.Vec3(lx, 0.0, 0.0)
            b = self.unit.Vec3(0.0, ly, 0.0)
            c = self.unit.Vec3(0.0, 0.0, lz)
            self.context.setPeriodicBoxVectors(a, b, c)
        self.context.setPositions(positions_nm * self.unit.nanometer)

    def energy_between(self, group_a, group_b):
        self.force.clearInteractionGroups()
        self.force.addInteractionGroup(set(group_a), set(group_b))
        self.force.updateParametersInContext(self.context)
        state = self.context.getState(getEnergy=True)
        return state.getPotentialEnergy().value_in_unit(self.unit.kilojoule_per_mole)


class _PmeEnergyEstimator:
    def __init__(self, openmm_system, cutoff, platform_name=None):
        mm, unit = _require_openmm()
        if not openmm_system.usesPeriodicBoundaryConditions():
            raise ValueError("PME requires periodic boundary conditions")
        self.system = mm.System()
        for i in range(openmm_system.getNumParticles()):
            self.system.addParticle(openmm_system.getParticleMass(i))
        nb = _find_nonbonded(openmm_system)
        self.force = mm.NonbondedForce()
        self.force.setNonbondedMethod(mm.NonbondedForce.PME)
        self.force.setCutoffDistance(cutoff * unit.nanometer)
        self.force.setUseDispersionCorrection(False)
        self.base_params = []
        for i in range(nb.getNumParticles()):
            charge, sigma, epsilon = nb.getParticleParameters(i)
            self.force.addParticle(charge, sigma, epsilon)
            self.base_params.append((charge, sigma, epsilon))
        self.base_exceptions = []
        for i in range(nb.getNumExceptions()):
            p1, p2, q, sig, eps = nb.getExceptionParameters(i)
            self.force.addException(p1, p2, q, sig, eps)
            self.base_exceptions.append((p1, p2, q, sig, eps))
        self.system.addForce(self.force)
        self.integrator = mm.VerletIntegrator(1.0 * unit.femtosecond)
        if platform_name:
            platform = mm.Platform.getPlatformByName(platform_name)
            self.context = mm.Context(self.system, self.integrator, platform)
        else:
            self.context = mm.Context(self.system, self.integrator)
        self.unit = unit
        self.active_mask = None

    def set_positions(self, positions_nm, box):
        a = self.unit.Vec3(box[0], 0.0, 0.0)
        b = self.unit.Vec3(0.0, box[1], 0.0)
        c = self.unit.Vec3(0.0, 0.0, box[2])
        self.context.setPeriodicBoxVectors(a, b, c)
        self.context.setPositions(positions_nm * self.unit.nanometer)

    def set_active_mask(self, active_mask):
        if self.active_mask is not None and np.array_equal(self.active_mask, active_mask):
            return
        self.active_mask = np.array(active_mask, dtype=bool)
        for i, (charge, sigma, epsilon) in enumerate(self.base_params):
            if self.active_mask[i]:
                self.force.setParticleParameters(i, charge, sigma, epsilon)
            else:
                self.force.setParticleParameters(i, charge * 0.0, sigma, epsilon * 0.0)
        for idx, (p1, p2, q, sig, eps) in enumerate(self.base_exceptions):
            if self.active_mask[p1] and self.active_mask[p2]:
                self.force.setExceptionParameters(idx, p1, p2, q, sig, eps)
            else:
                self.force.setExceptionParameters(idx, p1, p2, 0.0, sig, 0.0)
        self.force.updateParametersInContext(self.context)

    def energy(self):
        state = self.context.getState(getEnergy=True)
        return state.getPotentialEnergy().value_in_unit(self.unit.kilojoule_per_mole)


def gist(
    traj,
    system,
    openmm_system,
    openmm_topology,
    config: Optional[GistConfig] = None,
    platform_name: Optional[str] = None,
):
    config = config or GistConfig()
    if config.grid_spacing <= 0:
        raise ValueError("grid_spacing must be > 0")
    energy_method = config.energy_method.lower()
    if energy_method not in ("direct", "pme"):
        raise ValueError("energy_method must be 'direct' or 'pme'")

    waters = _water_groups(openmm_topology, config.water_resnames)
    if not waters:
        raise ValueError("no water residues found for GIST")
    water_atoms = [w[0] for w in waters]
    water_oxygens = [w[1] for w in waters]
    water_hydrogens = [w[2] for w in waters]

    all_water_atoms = set()
    for atoms in water_atoms:
        all_water_atoms.update(atoms)
    if config.solute_selection is not None:
        if system is None:
            raise ValueError("system is required when solute_selection is provided")
        solute_sel = system.select(config.solute_selection)
        solute_atoms = [int(i) for i in solute_sel.indices]
        if not solute_atoms:
            raise ValueError("solute_selection resolved to empty selection")
    else:
        solute_atoms = [i for i in range(openmm_system.getNumParticles()) if i not in all_water_atoms]
    water_other_atoms = []
    for atoms in water_atoms:
        others = [a for a in all_water_atoms if a not in atoms]
        water_other_atoms.append(others)

    max_chunk = config.chunk_frames or 128
    first = traj.read_chunk(max_chunk)
    if first is None:
        raise ValueError("trajectory has no frames")

    coords = np.asarray(first["coords"], dtype=np.float64) * config.length_scale
    if solute_atoms:
        origin, dims = _grid_from_positions(coords[0, solute_atoms, :], config.grid_spacing, config.padding)
    else:
        origin, dims = _grid_from_positions(coords[0, :, :], config.grid_spacing, config.padding)

    counts = np.zeros(tuple(dims), dtype=np.int64)
    energy_sw = np.zeros(tuple(dims), dtype=np.float64)
    energy_ww = np.zeros(tuple(dims), dtype=np.float64)
    frame_energy_sw = np.zeros(tuple(dims), dtype=np.float64)
    frame_energy_ww = np.zeros(tuple(dims), dtype=np.float64)
    orient_bins = max(int(config.orientation_bins), 1)
    orient_counts = np.zeros(tuple(dims) + (orient_bins,), dtype=np.int64)

    partitioner = _EnergyPartitioner(openmm_system, config.cutoff, platform_name=platform_name)
    pme_estimator = None
    if energy_method == "pme":
        pme_estimator = _PmeEnergyEstimator(openmm_system, config.cutoff, platform_name=platform_name)
        n_particles = openmm_system.getNumParticles()
        active_all = np.ones(n_particles, dtype=bool)
        active_solute = np.zeros(n_particles, dtype=bool)
        if solute_atoms:
            active_solute[solute_atoms] = True
        active_water = np.zeros(n_particles, dtype=bool)
        if all_water_atoms:
            active_water[list(all_water_atoms)] = True
    n_frames = 0

    if config.frame_indices is not None:
        frame_set = {int(i) for i in config.frame_indices if int(i) >= 0}
        max_frame_idx = max(frame_set) if frame_set else -1
    else:
        frame_set = None
        max_frame_idx = None

    chunk = first
    global_frame = 0
    while chunk is not None:
        coords = np.asarray(chunk["coords"], dtype=np.float64) * config.length_scale
        box = chunk.get("box")
        if box is not None:
            box = np.asarray(box, dtype=np.float64) * config.length_scale
        frames = coords.shape[0]
        for f in range(frames):
            if frame_set is not None and global_frame not in frame_set:
                global_frame += 1
                continue
            pos = coords[f]
            frame_box = box[f] if box is not None else None
            partitioner.set_positions(pos, box=frame_box)
            if pme_estimator is not None:
                if frame_box is None:
                    raise ValueError("PME requires periodic box")
                pme_estimator.set_positions(pos, frame_box)
            solute_center = pos[solute_atoms].mean(axis=0) if solute_atoms else pos.mean(axis=0)
            frame_energy_sw.fill(0.0)
            frame_energy_ww.fill(0.0)
            direct_sw_total = 0.0
            direct_ww_total = 0.0

            for w_idx, atoms in enumerate(water_atoms):
                oxy = water_oxygens[w_idx]
                if oxy is None:
                    continue
                voxel = _voxel_index(pos[oxy], origin, config.grid_spacing, dims)
                if voxel is None:
                    continue
                counts[voxel] += 1

                e_sw = partitioner.energy_between(solute_atoms, atoms) if solute_atoms else 0.0
                others = water_other_atoms[w_idx]
                e_ww = 0.5 * partitioner.energy_between(atoms, others) if others else 0.0
                frame_energy_sw[voxel] += e_sw
                frame_energy_ww[voxel] += e_ww
                direct_sw_total += e_sw
                direct_ww_total += e_ww

                hydrogens = water_hydrogens[w_idx]
                if len(hydrogens) >= 1:
                    dipole = pos[hydrogens].mean(axis=0) - pos[oxy]
                    rvec = pos[oxy] - solute_center
                    dipole_norm = np.linalg.norm(dipole)
                    r_norm = np.linalg.norm(rvec)
                    if dipole_norm > 0 and r_norm > 0:
                        cos_t = float(np.dot(dipole, rvec) / (dipole_norm * r_norm))
                        cos_t = max(-1.0, min(1.0, cos_t))
                        bin_idx = int(((cos_t + 1.0) * 0.5) * orient_bins)
                        if bin_idx == orient_bins:
                            bin_idx = orient_bins - 1
                        orient_counts[voxel + (bin_idx,)] += 1

            if pme_estimator is not None:
                pme_estimator.set_active_mask(active_all)
                e_all = pme_estimator.energy()
                pme_estimator.set_active_mask(active_solute)
                e_ss = pme_estimator.energy()
                pme_estimator.set_active_mask(active_water)
                e_ww = pme_estimator.energy()
                e_sw = e_all - e_ss - e_ww

                sw_scale = (e_sw / direct_sw_total) if direct_sw_total != 0.0 else 0.0
                ww_scale = (e_ww / direct_ww_total) if direct_ww_total != 0.0 else 0.0
                energy_sw += frame_energy_sw * sw_scale
                energy_ww += frame_energy_ww * ww_scale
            else:
                energy_sw += frame_energy_sw
                energy_ww += frame_energy_ww

            n_frames += 1
            global_frame += 1
            if config.max_frames is not None and n_frames >= config.max_frames:
                break
        if config.max_frames is not None and n_frames >= config.max_frames:
            break
        if max_frame_idx is not None and global_frame > max_frame_idx:
            break
        chunk = traj.read_chunk(max_chunk)

    voxel_volume = config.grid_spacing ** 3
    density = counts / (max(n_frames, 1) * voxel_volume)
    bulk_density = (
        float(config.bulk_density)
        if config.bulk_density is not None
        else counts.sum() / (max(n_frames, 1) * voxel_volume * density.size)
    )
    ratio = np.clip(density / max(bulk_density, 1e-12), 1e-12, None)
    ts_trans = -K_BOLTZ * config.temperature * density * np.log(ratio)

    ts_orient = np.zeros_like(density)
    with np.errstate(divide="ignore", invalid="ignore"):
        for idx in np.ndindex(*dims):
            total = counts[idx]
            if total <= 0:
                continue
            probs = orient_counts[idx] / total
            probs = probs[probs > 0]
            ts_orient[idx] = -K_BOLTZ * config.temperature * float(np.sum(probs * np.log(probs)))

    energy_sw = np.where(counts > 0, energy_sw / counts, 0.0)
    energy_ww = np.where(counts > 0, energy_ww / counts, 0.0)

    return GistResult(
        density=density,
        energy_sw=energy_sw,
        energy_ww=energy_ww,
        ts_trans=ts_trans,
        ts_orient=ts_orient,
        counts=counts,
        origin=origin,
        spacing=config.grid_spacing,
        n_frames=n_frames,
        voxel_volume=voxel_volume,
        bulk_density=bulk_density,
    )
