"""
LJTorch: A PyTorch-accelerated Lennard-Jones Calculator for ASE.

This module provides a GPU-capable calculator for multi-species Lennard-Jones
interactions with Periodic Boundary Conditions (PBC).
"""
from __future__ import annotations

import torch
import numpy as np
from typing import Dict, Tuple, Optional, Any, List
from ase.calculators.calculator import Calculator, all_changes
from ase.geometry import cell_to_cellpar

# Re-use the parameter table from the standard LJ module if possible, 
# otherwise define a local subset.
from .LJ import _DEFAULT_LJ_TABLE, _SIXTH_ROOT_OF_2, _get_vdw_radius, _infer_sigma_from_vdw, _infer_eps_from_size

class LJTorch(Calculator):
    """
    Experimental PyTorch-based Lennard-Jones calculator.
    Designed for small to medium systems where O(N^2) memory is acceptable 
    but GPU speed is desired.
    """
    implemented_properties = ["energy", "forces", "stress"]

    def __init__(self,
                 device: Optional[str] = None,
                 dtype: Optional[torch.dtype] = None,
                 species_params: Optional[Dict[str, Dict[str, float]]] = None,
                 pair_params: Optional[Dict[str, Dict[str, float]]] = None,
                 rc: Optional[float] = None,
                 shift: bool = True,
                 eps0: float = 0.010,
                 combination_rule: str = "LB",
                 **kwargs):
        super().__init__(**kwargs)
        
        # 1. Automatic device selection (Metal/MPS for Mac, ROCm for AMD on Linux)
        if device is None:
            if torch.cuda.is_available():
                # Note: PyTorch with ROCm (for AMD GPUs on Linux) 
                # uses "cuda" as an alias, so it will fall in here.
                device = "cuda"
            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
                device = "mps"
            elif hasattr(torch, "is_vulkan_available") and torch.is_vulkan_available():
                # Experimental Vulkan support (can work on AMD APUs)
                device = "vulkan"
            else:
                # Attempt to load DirectML for Windows (very useful for AMD APUs)
                try:
                    import torch_directml
                    if torch_directml.is_available():
                        device = torch_directml.device()
                    else:
                        device = "cpu"
                except ImportError:
                    device = "cpu"
        
        self.device = torch.device(device)

        # 2. Precision selection
        if dtype is None:
            # Note: Many MPS environments only support float32 or have limited float64 support.
            if self.device.type == "mps":
                dtype = torch.float32
            else:
                dtype = torch.float64
        
        self.dtype = dtype
        self.species_params = species_params or {}
        self.pair_params = pair_params or {}
        self.rc = rc
        self.shift = shift
        self.eps0 = eps0
        self.combination_rule = combination_rule

    def _resolve_parameters(self, symbols: List[str]) -> Tuple[torch.Tensor, torch.Tensor]:
        """Resolves Sigma and Epsilon matrices for the current configuration."""
        uniq = sorted(set(symbols))
        idx_map = {s: i for i, s in enumerate(uniq)}
        m = len(uniq)

        # 1. Resolve element-wise params
        per_elem = {}
        for s in uniq:
            if s in self.species_params:
                p = self.species_params[s]
                per_elem[s] = (float(p["sigma"]), float(p["epsilon"]))
            elif s in _DEFAULT_LJ_TABLE:
                per_elem[s] = _DEFAULT_LJ_TABLE[s]
            else:
                r = _get_vdw_radius(s)
                sig = _infer_sigma_from_vdw(r)
                eps = _infer_eps_from_size(r, self.eps0)
                per_elem[s] = (sig, eps)

        # 2. Build Matrices
        sigma_mat = torch.zeros((m, m), device=self.device, dtype=self.dtype)
        eps_mat = torch.zeros((m, m), device=self.device, dtype=self.dtype)

        for i, si in enumerate(uniq):
            for j, sj in enumerate(uniq):
                # Pair overrides
                key1, key2 = f"{si}-{sj}", f"{sj}-{si}"
                if key1 in self.pair_params:
                    p = self.pair_params[key1]
                    sij, eij = float(p["sigma"]), float(p["epsilon"])
                elif key2 in self.pair_params:
                    p = self.pair_params[key2]
                    sij, eij = float(p["sigma"]), float(p["epsilon"])
                else:
                    # Mixing
                    p1, p2 = per_elem[si], per_elem[sj]
                    if self.combination_rule == "geometric":
                        sij = np.sqrt(p1[0] * p2[0])
                        eij = np.sqrt(p1[1] * p2[1])
                    else: # LB
                        sij = 0.5 * (p1[0] + p2[0])
                        eij = np.sqrt(p1[1] * p2[1])
                
                sigma_mat[i, j] = sij
                eps_mat[i, j] = eij
        
        return sigma_mat, eps_mat

    def calculate(self, atoms=None, properties=('energy',), system_changes=all_changes):
        super().calculate(atoms, properties, system_changes)

        # 1. Transfer data to device
        # Using no_grad because we will compute forces and virial manually for performance
        with torch.no_grad():
            pos = torch.tensor(self.atoms.get_positions(), device=self.device, dtype=self.dtype)
            symbols = self.atoms.get_chemical_symbols()
            cell = torch.tensor(self.atoms.get_cell().array, device=self.device, dtype=self.dtype)
            
            n = pos.shape[0]
            if n == 0:
                self.results["energy"] = 0.0
                self.results["forces"] = np.zeros((0, 3))
                self.results["stress"] = np.zeros(6)
                return

            # 2. Resolve parameters
            sigma_mat, eps_mat = self._resolve_parameters(symbols)
            uniq = sorted(set(symbols))
            idx_map = {s: i for i, s in enumerate(uniq)}
            types = torch.tensor([idx_map[s] for s in symbols], device=self.device, dtype=torch.long)

            i_types = types.unsqueeze(1)
            j_types = types.unsqueeze(0)
            sig = sigma_mat[i_types, j_types]
            eps = eps_mat[i_types, j_types]

            # 3. Pair distances with MIC
            diff = pos.unsqueeze(1) - pos.unsqueeze(0)
            is_orthogonal = torch.allclose(cell, torch.diag(torch.diag(cell)))
            
            if is_orthogonal:
                L = torch.diag(cell)
                # Check for zero size in case of non-PBC
                L_eff = L.clone()
                L_eff[L_eff < 1e-9] = 1e9 
                diff = diff - L_eff * torch.round(diff / L_eff)
            else:
                inv_cell = torch.inverse(cell)
                frac_diff = diff @ inv_cell
                frac_diff = frac_diff - torch.round(frac_diff)
                diff = frac_diff @ cell

            r2 = torch.sum(diff**2, dim=-1)
            
            # Mask: i < j and rc
            mask = torch.triu(torch.ones((n, n), device=self.device, dtype=torch.bool), diagonal=1)
            if self.rc is not None:
                 rc2 = self.rc ** 2
                 mask = mask & (r2 < rc2)
                 
            r2_masked = r2 + (~mask).to(self.dtype)
            
            # 4. Potentials
            sr2 = (sig**2) / r2_masked
            sr6 = sr2**3
            sr12 = sr6**2
            
            # Energy
            v_ij = 4.0 * eps * (sr12 - sr6)
            if self.rc is not None and self.shift:
                sr2c = (sig / self.rc)**2
                sr6c = sr2c**3
                v_rc = 4.0 * eps * (sr6c**2 - sr6c)
                v_ij = (v_ij - v_rc)
            
            energy = torch.sum(v_ij[mask])

            # 5. Forces (Manual derivative)
            # F_ij = -dV/dr * (r_ij / |r_ij|) = 24*eps*(2*sr12 - sr6)/r2 * r_ij
            coef = 24.0 * eps * (2.0 * sr12 - sr6) / r2_masked
            coef = coef * mask.to(self.dtype)
            
            # fij: (N, N, 3)
            f_ij_vec = coef.unsqueeze(-1) * diff
            
            # sum over i
            # force on i is sum_{j} f_ij - sum_{j} f_ji (but mask i < j)
            forces = torch.sum(f_ij_vec, dim=1) - torch.sum(f_ij_vec, dim=0)

            # 6. Stress (Virial)
            # Pick only unique pairs to avoid double counting if we didn't use i < j
            pair_idx = torch.where(mask)
            r_pairs = diff[pair_idx]
            f_pairs = f_ij_vec[pair_idx]
            virial = torch.mm(r_pairs.t(), f_pairs) # (3, 3)
            
            volume = torch.abs(torch.det(cell))
            stress_voigt = np.zeros(6)
            if volume > 1e-12:
                stress_tensor = -virial / volume
                st = stress_tensor.cpu().numpy()
                stress_voigt = np.array([st[0,0], st[1,1], st[2,2], st[1,2], st[0,2], st[0,1]])

            # Transfer results back to CPU
            self.results["energy"] = float(energy.item())
            self.results["forces"] = forces.cpu().numpy()
            self.results["stress"] = stress_voigt

def lj_torch(**kwargs) -> LJTorch:
    """Factory function to instantiate LJTorch."""
    return LJTorch(**kwargs)
