import time
import math
import numpy as np
from scipy.spatial import cKDTree
from numba import njit, prange
from ase.data import atomic_masses
from aegon.libutils import sort_by_energy

#--- USR ----------------------------------------------------------------------
def compute_usr(lista_atoms, mode="mono", dtype=np.float32):
    pos_cat, masses, offsets = pack_atoms_list(lista_atoms)
    M = offsets.shape[0] - 1

    if mode == "mono":
        X = np.empty((M, 8), dtype=np.float64)
        usr_batch12(pos_cat, offsets, X)
    else:
        X = np.empty((M, 16), dtype=np.float64)
        usr_batch24(pos_cat, masses, offsets, X)

    return X.astype(dtype, copy=False)


def pack_atoms_list(lista_atoms):
    M = len(lista_atoms)
    counts  = np.fromiter((len(a) for a in lista_atoms), dtype=np.int64, count=M)
    offsets = np.zeros(M + 1, dtype=np.int64)
    np.cumsum(counts, out=offsets[1:])
    totalN = int(offsets[-1])

    pos_cat = np.empty((totalN, 3), dtype=np.float64)
    masses  = np.empty(totalN, dtype=np.float64)

    p = 0
    for at in lista_atoms:
        n = len(at)
        pos_cat[p:p+n, :] = at.get_positions()
        Z = at.get_atomic_numbers()
        masses[p:p+n] = atomic_masses[Z]
        p += n

    return pos_cat, masses, offsets


@njit(fastmath=True)
def four_points(pos):
    n = pos.shape[0]

    ctd0 = 0.0; ctd1 = 0.0; ctd2 = 0.0
    for i in range(n):
        ctd0 += pos[i, 0]
        ctd1 += pos[i, 1]
        ctd2 += pos[i, 2]
    invn = 1.0 / n
    ctd0 = ctd0 * invn; ctd1 = ctd1 * invn; ctd2 = ctd2 * invn

    i_cst = 0; i_fct = 0
    dmin = math.inf
    dmax = -math.inf
    for i in range(n):
        dx = pos[i, 0] - ctd0
        dy = pos[i, 1] - ctd1
        dz = pos[i, 2] - ctd2
        di = math.sqrt(dx*dx + dy*dy + dz*dz)
        if di < dmin:
            dmin = di; i_cst = i
        if di > dmax:
            dmax = di; i_fct = i

    cst0, cst1, cst2 = pos[i_cst, 0], pos[i_cst, 1], pos[i_cst, 2]
    fct0, fct1, fct2 = pos[i_fct, 0], pos[i_fct, 1], pos[i_fct, 2]

    i_ftf = 0
    dmax2 = -math.inf
    for i in range(n):
        dx = pos[i, 0] - fct0
        dy = pos[i, 1] - fct1
        dz = pos[i, 2] - fct2
        di = math.sqrt(dx*dx + dy*dy + dz*dz)
        if di > dmax2:
            dmax2 = di; i_ftf = i

    ftf0, ftf1, ftf2 = pos[i_ftf, 0], pos[i_ftf, 1], pos[i_ftf, 2]
    return (ctd0, ctd1, ctd2,
            cst0, cst1, cst2,
            fct0, fct1, fct2,
            ftf0, ftf1, ftf2)


@njit(fastmath=True)
def moments_point(pos, p0, p1, p2):
    n = pos.shape[0]

    s1 = 0.0
    for i in range(n):
        dx = pos[i,0]-p0; dy = pos[i,1]-p1; dz = pos[i,2]-p2
        d  = math.sqrt(dx*dx + dy*dy + dz*dz)
        s1 += d
    mu = s1 / n

    m2 = 0.0; m3 = 0.0
    for i in range(n):
        dx = pos[i,0]-p0; dy = pos[i,1]-p1; dz = pos[i,2]-p2
        d  = math.sqrt(dx*dx + dy*dy + dz*dz)
        r  = d - mu
        r2 = r*r
        m2 += r2
        m3 += r2*r

    return m2/n, m3/n


@njit(fastmath=True)
def moments_point_weighted(pos, p0, p1, p2, w):
    n = pos.shape[0]

    s1 = 0.0
    for i in range(n):
        dx = pos[i,0]-p0; dy = pos[i,1]-p1; dz = pos[i,2]-p2
        d  = math.sqrt(dx*dx + dy*dy + dz*dz) * w[i]
        s1 += d
    mu = s1 / n

    m2 = 0.0; m3 = 0.0
    for i in range(n):
        dx = pos[i,0]-p0; dy = pos[i,1]-p1; dz = pos[i,2]-p2
        d  = math.sqrt(dx*dx + dy*dy + dz*dz) * w[i]
        r  = d - mu
        r2 = r*r
        m2 += r2
        m3 += r2*r

    return m2/n, m3/n


@njit(fastmath=True)
def usr_mono(pos):
    ctd0,ctd1,ctd2, cst0,cst1,cst2, fct0,fct1,fct2, ftf0,ftf1,ftf2 = four_points(pos)

    r1 = moments_point(pos, ctd0,ctd1,ctd2)
    r2 = moments_point(pos, cst0,cst1,cst2)
    r3 = moments_point(pos, fct0,fct1,fct2)
    r4 = moments_point(pos, ftf0,ftf1,ftf2)

    out = np.empty(8, dtype=np.float64)
    out[0:2] = r1; out[2:4] = r2; out[4:6] = r3; out[6:8] = r4
    return out


@njit(fastmath=True)
def usr_multi(pos, w):
    ctd0,ctd1,ctd2, cst0,cst1,cst2, fct0,fct1,fct2, ftf0,ftf1,ftf2 = four_points(pos)

    r1 = moments_point(pos, ctd0,ctd1,ctd2)
    r2 = moments_point(pos, cst0,cst1,cst2)
    r3 = moments_point(pos, fct0,fct1,fct2)
    r4 = moments_point(pos, ftf0,ftf1,ftf2)

    rw1 = moments_point_weighted(pos, ctd0,ctd1,ctd2, w)
    rw2 = moments_point_weighted(pos, cst0,cst1,cst2, w)
    rw3 = moments_point_weighted(pos, fct0,fct1,fct2, w)
    rw4 = moments_point_weighted(pos, ftf0,ftf1,ftf2, w)

    out = np.empty(16, dtype=np.float64)
    out[0:2]   = r1;  out[2:4]   = r2;  out[4:6]   = r3;  out[6:8]   = r4
    out[8:10]  = rw1; out[10:12] = rw2; out[12:14] = rw3; out[14:16] = rw4
    return out


@njit(parallel=True, fastmath=True)
def usr_batch12(pos_cat, offsets, out8):
    M = offsets.shape[0] - 1
    for m in prange(M):
        a = offsets[m]; b = offsets[m+1]
        pos = pos_cat[a:b, :]
        out8[m, :] = usr_mono(pos)


@njit(parallel=True, fastmath=True)
def usr_batch24(pos_cat, masses, offsets, out16):
    M = offsets.shape[0] - 1
    for m in prange(M):
        a = offsets[m]; b = offsets[m+1]
        pos = pos_cat[a:b, :]
        w_raw = masses[a:b]

        s = 0.0
        for i in range(w_raw.size):
            s += w_raw[i]
        invavg = w_raw.size / s

        w = np.empty(w_raw.size, dtype=np.float64)
        for i in range(w_raw.size):
            w[i] = w_raw[i] * invavg

        out16[m, :] = usr_multi(pos, w)

#--- Filtros ------------------------------------------------------------------
def filtro_est(lista_atoms, sim=0.97, dE_max=None, mode="mono", dtype=np.float32):
    t1 = time.perf_counter()

    lista_sorted = sort_by_energy(lista_atoms)
    E = np.array([float(a.info["e"]) for a in lista_sorted], dtype=np.float64)

    X = compute_usr(lista_sorted, mode=mode, dtype=dtype)
    N, D = X.shape

    dmax = D * (1.0 - sim) / sim

    keep1 = prefiltro(X, E, dmax, dE_max)
    X1 = X[keep1]
    E1 = E[keep1]
    lista1 = [lista_sorted[i] for i in np.nonzero(keep1)[0]]

    keep2 = filtro_kdtree(X1, E1, dmax, dE_max)
    lista_out = [lista1[i] for i in np.nonzero(keep2)[0]]

    t2 = time.perf_counter()
    print(f"Filtro_USR: {t2 - t1:.3f} s | in={len(lista_atoms)} out={len(lista_out)}")

    return lista_out

def filtro_est_ref_search(ref_atoms, search_atoms, sim=0.97, dE_max=None, mode="mono", dtype=np.float32):
    X_ref = compute_usr(ref_atoms,   mode=mode, dtype=dtype)
    X_sea = compute_usr(search_atoms, mode=mode, dtype=dtype)

    D = X_ref.shape[1]
    dmax = D * (1.0 - sim) / sim

    if dE_max is None:
        keep = filtro_kdtree_ref_search(X_ref, None, X_sea, None, dmax, None)
    else:
        E_ref = np.array([float(a.info["e"]) for a in ref_atoms], dtype=np.float64)
        E_sea = np.array([float(a.info["e"]) for a in search_atoms], dtype=np.float64)
        keep = filtro_kdtree_ref_search(X_ref, E_ref, X_sea, E_sea, dmax, dE_max)

    search_out = [a for a, k in zip(search_atoms, keep) if k]
    return search_out

def prefiltro(X, E, dmax, dE_max):
    X = np.asarray(X, dtype=np.float32, order="C")
    E = np.asarray(E, dtype=np.float64)

    if dE_max is None:
        return prefiltro_core(X, dmax)
    else:
        return prefiltro_energy_core(X, E, dmax, float(dE_max))

def filtro_kdtree(X, E, dmax, dE_max):
    X = np.asarray(X, dtype=np.float32, order="C")
    E = np.asarray(E, dtype=np.float64)
    N = X.shape[0]

    tree = cKDTree(X)
    keep = np.ones(N, dtype=bool)

    if dE_max is None:
        for i in range(N):
            if not keep[i]:
                continue
            neigh = tree.query_ball_point(X[i], r=dmax, p=1)
            for j in neigh:
                if j > i:
                    keep[j] = False
    else:
        dEmax = float(dE_max)
        for i in range(N):
            if not keep[i]:
                continue
            Ei = E[i]
            neigh = tree.query_ball_point(X[i], r=dmax, p=1)
            for j in neigh:
                if j <= i:
                    continue
                dE = E[j] - Ei
                if dE < 0.0:
                    dE = -dE
                if dE <= dEmax:
                    keep[j] = False
    return keep

def filtro_kdtree_ref_search(X_ref, E_ref, X_search, E_search, dmax, dE_max=None):
    X_ref    = np.asarray(X_ref,    dtype=np.float32, order="C")
    X_search = np.asarray(X_search, dtype=np.float32, order="C")

    if dE_max is not None:
        E_ref    = np.asarray(E_ref,    dtype=np.float64)
        E_search = np.asarray(E_search, dtype=np.float64)

    tree = cKDTree(X_ref)
    M = X_search.shape[0]
    keep = np.ones(M, dtype=bool)

    if dE_max is None:
        for i in range(M):
            neigh = tree.query_ball_point(X_search[i], r=dmax, p=1)
            if len(neigh) > 0:
                keep[i] = False
    else:
        dEmax = float(dE_max)
        for i in range(M):
            neigh = tree.query_ball_point(X_search[i], r=dmax, p=1)
            if len(neigh) == 0:
                continue
            Ei = E_search[i]
            dup = False
            for j in neigh:
                dE = E_ref[j] - Ei
                if dE < 0.0:
                    dE = -dE
                if dE <= dEmax:
                    dup = True
                    break
            if dup:
                keep[i] = False

    return keep

@njit
def manhattan_threshold(x, y, dmax):
    s = 0.0
    D = x.shape[0]
    for d in range(D):
        diff = x[d] - y[d]
        if diff < 0.0:
            diff = -diff
        s += diff
        if s > dmax:
            return False
    return True

@njit
def prefiltro_core(X, dmax):
    N, D = X.shape
    keep = np.ones(N, dtype=np.bool_)
    last_keep = 0
    for i in range(1, N):
        if manhattan_threshold(X[i], X[last_keep], dmax):
            keep[i] = False
        else:
            last_keep = i
    return keep

@njit
def prefiltro_energy_core(X, E, dmax, dE_max):
    N, D = X.shape
    keep = np.ones(N, dtype=np.bool_)
    last_keep = 0

    for i in range(1, N):
        dE = E[i] - E[last_keep]
        if dE < 0.0:
            dE = -dE
        if dE > dE_max:
            last_keep = i
            continue

        if manhattan_threshold(X[i], X[last_keep], dmax):
            keep[i] = False
        else:
            last_keep = i
    return keep
