import jax
import jax.numpy as jnp
import jax.scipy as jsp
from jaxfmm.basis import inv_mpl_idx, eval_singular_basis
from time import perf_counter
import numpy as np
import pyvista as pv
import os
from functools import partial

__all__ = ["gen_stats", "gen_hierarchy_vtk", "gen_wellsep_vtk", "gen_pts_vtk", "time_function"]

def gen_stats(pts, eval_pts, idcs, boxcenters, eval_boxcenters, mpl_cnct, dir_cnct, img_cnct, lvl_info, s, periodic_axes, p, theta, print_stats = False, **kwargs):
    """Generate hierarchy stats. Use the hierarchy information from gen_hierarchy as input."""
    stats = {"p": p,
             "theta": theta,
             "max_l": len(boxcenters)-1,
             "eval_max_l": len(eval_boxcenters)-1,
             "num_pts": pts.shape[0],
             "num_eval_pts": eval_pts.shape[0],
             "num_child": 2**s,
             "pts_per_box": idcs[2][0].shape[1],
             "eval_pts_per_box": idcs[3][0].shape[1]}
    periodic = len(periodic_axes)>0
    mem_hierarch = idcs[0][0].nbytes + idcs[0][1].nbytes
    if(idcs[1] is not idcs[0]):
        mem_hierarch += idcs[1][0].nbytes + idcs[1][1].nbytes
    if(idcs[2] is not idcs[0]):
        mem_hierarch += idcs[2][0].nbytes + idcs[2][1].nbytes
    if((idcs[3] is not idcs[1]) and (idcs[3] is not idcs[2])):
        mem_hierarch += idcs[3][0].nbytes + idcs[3][1].nbytes
    mem_hierarch = dir_cnct.nbytes + img_cnct.nbytes
    M2L_nums = []
    M2L_pad_fracs = []
    n_img = img_cnct.shape[0]

    for i in range(len(lvl_info)-1):
        _, l_src = lvl_info[i]
        num_nopad = (mpl_cnct[i] < boxcenters[l_src].shape[0]*n_img).sum().tolist()
        M2L_nums.append(mpl_cnct[i].size)
        mem_hierarch += mpl_cnct[i].nbytes
        M2L_pad_fracs.append(1 if num_nopad==0 else M2L_nums[-1] / num_nopad)

    for i in range(stats["max_l"]):
        mem_hierarch += boxcenters[i].nbytes
        if 'boxlens' in kwargs:
            mem_hierarch += kwargs.get("boxlens")[i].nbytes

    if(eval_pts is not pts):
        for i in range(stats["eval_max_l"]):
            mem_hierarch += eval_boxcenters[i].nbytes
            if 'eval_boxlens' in kwargs:
                mem_hierarch += kwargs.get("eval_boxlens")[i].nbytes

    stats["M2L_nums"] = M2L_nums
    stats["lvl_info"] = lvl_info
    stats["M2L_pad_fracs"] = M2L_pad_fracs
    stats["direct_num"] = dir_cnct.size
    dir_no_pad = (dir_cnct < boxcenters[lvl_info[-1][1]].shape[0]*n_img).sum().tolist()
    stats["direct_pad_frac"] = (dir_cnct.size / dir_no_pad) if dir_no_pad != 0 else 1
    stats["compression_ratio"] = dir_no_pad * idcs[3][0].shape[1] * idcs[2][0].shape[1] / (pts.shape[0]*eval_pts.shape[0])
    stats["mem_hierarch"] = mem_hierarch
    stats["mem_coeffs"] = 4*2*((p+1)**2)*(8**(stats["max_l"]+1)-1)/7 # 4 bytes, 2 arrays, (p+1)**2 coeffs and nboxes
    stats["mem_pts_chrgs"] = (pts.nbytes*4)/3
    if(eval_pts is not pts):
        stats["mem_pts_chrgs"] += eval_pts.nbytes
    if(print_stats):
        print("---------------------FMM Hierarchy Stats---------------------")
        print("p = %i, theta = %.2f, %i children per box"%(p,theta,stats["num_child"]))
        if(periodic): 
            print("Periodic boundary on axes: ", periodic_axes)
        print("%i points, %i levels, %i particles per box"%(stats["num_pts"],stats["max_l"], stats["pts_per_box"]))
        if(eval_pts is not pts):
            print("%i eval points, %i eval levels, %i eval points per box"%(stats["num_eval_pts"],stats["eval_max_l"],stats["eval_pts_per_box"]))
        print("")
        print(" (l_eval, l_src) | padding frac. | interactions")
        print("-----------------------------------------------")
        for i in range(len(lvl_info)-1):
            l_eval, l_src = lvl_info[i]
            print(" M2L (%2i,%2i)     |     %5.2f     | %11i"%(l_eval, l_src, stats["M2L_pad_fracs"][i], stats["M2L_nums"][i]))
        totsum = jnp.array(stats["M2L_nums"]).sum()
        print("-----------------------------------------------")
        print(" M2L Total:            %5.2f       %11i"%(totsum/jnp.array(stats["M2L_nums"]).dot(1/jnp.array(stats["M2L_pad_fracs"])),totsum))
        print("-----------------------------------------------")
        print(" dir (%2i,%2i)     |     %5.2f     | %11i"%(lvl_info[-1][0], lvl_info[-1][1], stats["direct_pad_frac"], stats["direct_num"]))
        print("-----------------------------------------------")
        print("\nDirect interaction compression (without padding): %.2e"%(stats["compression_ratio"]))
        print("\nMemory used by the hierarchy:            %.2e Bytes"%stats["mem_hierarch"])
        print("Memory used by mpl + local coeffs (p=%i): %.2e Bytes"%(p, stats["mem_coeffs"])) 
        print("Memory used by the points + charges:     %.2e Bytes"%stats["mem_pts_chrgs"])
        print("-------------------------------------------------------------")
    return stats

def make_pyvista_mesh(boxcenters, boxlens):
    r"""
    Generate pyvista mesh of boxes with centers at boxcenters and diagonals boxlens.
    """
    pv_shifts = np.array([[-1,-1,-1],[1,-1,-1],[1,1,-1],[-1,1,-1],[-1,-1,1],[1,-1,1],[1,1,1],[-1,1,1]])
    pts = np.zeros((boxcenters.shape[0],8,3))
    for i in range(pv_shifts.shape[0]):
        pts[:,i,:] = boxcenters + boxlens/2 * pv_shifts[i,None,:]
    pts = pts.reshape((-1,3))
    cells = np.arange(pts.shape[0],dtype=np.int32).reshape((-1,8))
    return pv.UnstructuredGrid({pv.CellType.HEXAHEDRON: cells}, pts)

def gen_hierarchy_vtk(boxcenters, boxlens, eval_boxcenters, eval_boxlens, dir="hierarchy", **kwargs):
    r"""
    Output a series of vtk files showing the FMM hierarchy on every level. Does not show virtual PBC images.
    """
    eval_max_l = len(eval_boxcenters)
    max_l = len(boxcenters)
    if not os.path.exists(dir):
        os.makedirs(dir)
    for l in range(max_l):
        mesh = make_pyvista_mesh(boxcenters[l],boxlens[l])
        mesh.cell_data["src_level"] = np.ones(mesh.points.shape[0]//8)*l
        mesh.save("%s/src_level_%i.vtk"%(dir,l))
    
    if(eval_boxcenters is not boxcenters or eval_boxlens is not boxlens):
        for l in range(eval_max_l):
            mesh = make_pyvista_mesh(eval_boxcenters[l],eval_boxlens[l])
            mesh.cell_data["eval_level"] = np.ones(mesh.points.shape[0]//8)*l
            mesh.save("%s/eval_level_%i.vtk"%(dir,l))

def gen_pts_vtk(pts, eval_pts, dir="hierarchy", **kwargs):
    r"""
    Output the hierarchy points as vtk.
    """
    if not os.path.exists(dir):
        os.makedirs(dir)
    point_cloud = pv.PolyData(np.array(pts))
    point_cloud.save("%s/src_pts.vtk"%(dir))
    point_cloud = pv.PolyData(np.array(eval_pts))
    point_cloud.save("%s/eval_pts.vtk"%(dir))

def gen_wellsep_vtk(id, boxcenters, boxlens, eval_boxcenters, eval_boxlens, mpl_cnct, dir_cnct, img_cnct, lvl_info, s, pbc_ws, dir="hierarchy", **kwargs):
    r"""
    Output a vtk file showing all the boxes that are considered for potential calculation, given an id of a box on the highest level. Includes virtual PBC images.
    """
    if not os.path.exists(dir):
        os.makedirs(dir)
    max_l_eval = lvl_info[-1][0]
    n_chi = 2**s
    periodic = len(pbc_ws)>0
    meshlist = []

    L = boxlens[0][0]
    for l, img_cnct_ws in enumerate(pbc_ws):
        mesh = make_pyvista_mesh(boxcenters[0] - img_cnct_ws[1],jnp.tile(L,(img_cnct_ws[1].shape[0],1)))
        mesh.cell_data["eval_level"] = np.ones(mesh.points.shape[0]//8)*(-l)
        mesh.cell_data["src_level"] = np.ones(mesh.points.shape[0]//8)*(-l)
        meshlist.append(mesh)
        L += jnp.ptp(img_cnct_ws[0],axis=0)

    for i in range(len(lvl_info)-1):
        l_eval, l_src = lvl_info[i]
        extract = mpl_cnct[i][id//(n_chi**(max_l_eval-l_eval))]
        n_src = n_chi**l_src
        if(periodic):
            extract = extract[extract < (n_src*img_cnct.shape[0])]
            mesh = make_pyvista_mesh(boxcenters[l_src][extract%n_src] + img_cnct[extract//n_src],boxlens[l_src][extract%n_src])
        else:
            extract = extract[extract < n_src]
            mesh = make_pyvista_mesh(boxcenters[l_src][extract],boxlens[l_src][extract])
        mesh.cell_data["eval_level"] = np.ones(mesh.points.shape[0]//8)*(l_eval)
        mesh.cell_data["src_level"] = np.ones(mesh.points.shape[0]//8)*(l_src)
        meshlist.append(mesh)

    l_eval, l_src = lvl_info[-1]
    extract = dir_cnct[id]
    n_src = n_chi**l_src
    if(periodic):
        extract = extract[extract<n_src*img_cnct.shape[0]]
        mesh = make_pyvista_mesh(boxcenters[-1][extract%n_src] + img_cnct[extract//n_src],boxlens[-1][extract%n_src])
    else:
        extract = extract[extract < n_src]
        mesh = make_pyvista_mesh(boxcenters[l_src][extract],boxlens[l_src][extract])

    mesh.cell_data["eval_level"] = np.ones(mesh.points.shape[0]//8)*(l_eval + 1)
    mesh.cell_data["src_level"] = np.ones(mesh.points.shape[0]//8)*(l_src + 1)
    meshlist.append(mesh)

    mesh = make_pyvista_mesh(eval_boxcenters[l_eval][id][None,:],eval_boxlens[l_eval][id][None,:]*1.001)   # the multiplication fixes pyvista getting confused for eval_pts = pts
    mesh.cell_data["eval_level"] = np.ones(1)*(l_eval + 2)
    mesh.cell_data["src_level"] = np.ones(1)*(l_src + 2)
    meshlist.append(mesh)

    mergedmesh = pv.merge(meshlist)
    mergedmesh.save("%s/wellsep_%i.vtk"%(dir,id))

@partial(jax.jit, static_argnames=['p'])
def eval_mpls(mpls, eval_pts, boxcenters, p):
    r"""
    Evaluate multipole expansions.
    """
    sing = eval_singular_basis(eval_pts - boxcenters[None,:],p)
    ms, _ = inv_mpl_idx(jnp.arange((p+1)**2))
    prefac = ((2-(ms==0))*(1-2*(ms<0)))[None,None,:]
    res = (mpls * sing * prefac).sum(axis=2)
    return res / (4*jnp.pi)

@partial(jax.jit, static_argnames=['p'])
def get_locs(padded_pts, padded_chrgs, boxcenters, p):
    r"""
    Obtain local expansions directly.
    """
    dist = padded_pts - boxcenters[:,None]
    sing = eval_singular_basis(dist,p)
    return (sing * padded_chrgs[...,None]).sum(axis=1)

def binom(x, y):
  return jnp.exp(jsp.special.gammaln(x + 1) - jsp.special.gammaln(y + 1) - jsp.special.gammaln(x - y + 1))

def gen_multipole_dist(m, n, eps = 0.5):
    r"""
    Generate a point charge distribution corresponding to a specific multipole moment (Majic, Matt. (2022). 
    Point charge representations of multipoles. European Journal of Physics. 43. 10.1088/1361-6404/ac578b.)
    """
    if(m == 0):   # axial
        k = jnp.arange(-n, n+1, 2)
        chrgs = (-1)**((n-k)/2) * binom(n, (n-k)/2.0) / (jsp.special.factorial(n) * (2*eps)**n)
        pts = jnp.zeros((k.shape[0],3))
        pts = pts.at[:,2].set(k*eps)
    else:         # (stacked) bracelet
        rotate = m < 0
        m = abs(m)      # we work with the real basis and rotate later
        knum = n-m+1
        jnum = 2*m
        j = jnp.tile(jnp.arange(jnum),knum)
        k = jnp.repeat(jnp.arange(-n+m,n-m+1,2),jnum)
        phi = (j-0.5) * jnp.pi/m if rotate else j * jnp.pi/m
        pts = jnp.array([eps*jnp.cos(phi), eps*jnp.sin(phi), k*eps]).T
        chrgs = 4**(m-1) * jsp.special.factorial(m-1) / ((2*eps)**n * jsp.special.factorial(n-m)) * (-1)**((n-m-k)/2 + j) * binom(n-m,(n-m-k)/2)
    return pts, chrgs

def time_function(func, nruns = 10, print_times = True, **kwargs):
    r"""
    Time function func, obtaining the best of nruns runs and optionally printing the results. Function args are supplied as additional kwargs.
    """
    best = jnp.inf
    for i in range(nruns):
        T0 = perf_counter()
        res = jax.block_until_ready(func(**kwargs))
        T1 = perf_counter()
        best = min(T1-T0,best)
        if(i==0):
            comp = best
    if(print_times):
        print("%s compilation: %.2e s"%(func.__name__,comp))
        print("%s runtime: %.2e s"%(func.__name__,best))
    return res, comp, best