import ebf
import os
import tempfile
import enlink
import matplotlib.pyplot as plt
import numpy as np

# This code shows how to get smoothing lengths from enlink.
# These smoothing lengths are used by Galaxia to sample N-body particles

# Installation and Usage for Pythin version of enlink
# tar -zxvf enlink-0.1.0.tar.gz
# cd enlink-0.1.0
# python setup.py install
# >>>import enlink
# >>>help(enlink)
# For latest ebfpy
# pip install --user --upgrade git+https://github.com/sanjibs/ebfpy.git@master


def enlink_galaxia(X,neighbors=64,density_method='sb'):
    """

    Intended for use with Galaxia.
    Computes density and smoothing length (metric) of given points.
    Galaxia uses cubic cells. For 3d density_method='' does it.
    For 6d phase space one has to use density_method='sb' for disc galaxies 
    and density_method='c' for halo.

    Parameters
    ----------
    X : {array-like} of shape (n_samples, n_features)
        Data points whose density is to be computed

    neighbors : int, optional (default=64)
        The number of neighbors to use for computing density

    density_method : str, optional (default='sbr')
        The default is the best possible metric calculation and is the
        recommended option. It computes an adaptive metric using an entropy
        based scheme. A bootstrap technique used for
        smoothing the metric. Metric can have non diagonal 
        elements also which is useful for anisotropic structures.
        For diagonal metric just use 'sb'.
        Other option is to set it to blank which will cluster using a simple
        Euclidean metric. See gmetric for more global options.
        Blank also enables space the EnBiD partitioning scheme without 
        rotation and bootstrap and with cubic cells.
        'c'  is used for cubic cells and is useful for dims like 6,9,12 etc.
 

    Returns
    -------
    density_estimate : array of shape (n_targets)

    """
    # write data to files

    fd,filename=tempfile.mkstemp(suffix='.ebf')
    tempstatus=True
    os.close(fd)
    ebf.write(filename,'/x',X,'w')

    gmetric=0
    custom_metric=''
    if X.ndim != 2:
        raise RuntimeError('X.ndim should be 2')
    else:
        dim=X.shape[1]
        
    # run Enlink
    s='Enlink -dm{0:s} --ngb={1:d} --gmetric={2:d} --dim={3:d} {4:s}'.format(density_method,neighbors,gmetric,dim,filename)
    enlink.run(s,custom_metric)

    d=ebf.read(filename[0:-4]+'_den.ebf','/')
    print(d.keys())
    # read data
    density=ebf.read(filename[0:-4]+'_den.ebf','/density')
    if ebf.containsKey(filename[0:-4]+'_den.ebf','/h_cubic'):
        h_smooth=ebf.read(filename[0:-4]+'_den.ebf','/h_cubic')
        h_smooth=np.vstack([h_smooth[:,0],h_smooth[:,0],h_smooth[:,0],h_smooth[:,1],h_smooth[:,1],h_smooth[:,1]]).transpose()*1.0
    else:
        h_smooth=ebf.read(filename[0:-4]+'_den.ebf','/h_smooth')
    # d=ebf.read(filename[0:-4]+'_den.ebf','/')
    print(h_smooth.shape)
    
    # ebf.info does not work in python3 yet will update ebf soon
    #ebf.info(filename[0:-4]+'_den.ebf')

    # delete temporary files
    fname=filename
    if os.path.exists(fname):
        os.remove(fname)

    fname=filename[0:-4]+'_den.ebf'
    if os.path.exists(fname):
        os.remove(fname)
            
    return density,h_smooth



# Note for disc galaxies density_method='sb' does better when 
# galaxy is aligned with x,y,z axes.

# Read Bullock and Johnston data,=. Here we read one satellelite 
# Modify this part as suited for your case.
# Basically for N points, data should be loaded into an array of shape [N,6]
pos=ebf.read('/work/sharma/GsynthData/nbody1/bj/sat0000.ebf','/pos3')
vel=ebf.read('/work/sharma/GsynthData/nbody1/bj/sat0000.ebf','/vel3')
# make the array have shape (N,6) phase space        
X=np.hstack([pos,vel])

# for 6d phase space diagonal metric
density,h_smooth=enlink_galaxia(X, density_method='sb')   
ebf.write('sat0000_d6n64_den.ebf','/h_smooth',h_smooth,'w')


# For cubic metric as in old Galaxia
# for 3d position space 
density,h_smooth=enlink_galaxia(X[:,0:3], density_method='')   
ebf.write('sat0000_d3n64_den.ebf','/h_cubic',h_smooth[:,0],'w')

# for 6d phase space
density,h_smooth=enlink_galaxia(X, density_method='c')   
h_cubic=np.vstack([h_smooth[:,0],h_smooth[:,3]]).transpose()*1.0
ebf.write('sat0000_d6n64_den.ebf','/h_cubic',h_cubic,'w')