import numpy as np
import matplotlib.pyplot as plt
from .. import nptipsyreader
from .. import readstat
import glob
import re
import pdb


def plot():
    box = 25.
    h = 0.6777
    darkmass = 339072.93827684096 #Msol
    
    mffilenames = glob.glob('*.mf') #'plk13.z10.mf'
    simfilename = 'out_0.list'
    
#### read in rockstar output ####
    f = open(simfilename)
    names = f.readline().strip('\n#').split(' ')
    dtype = np.concatenate((['int', 'int'], (36*'float ').split(' ')[:-1]))
    f.close()
    sim = np.genfromtxt(simfilename, dtype=dtype, names=names, skiprows=1)
    
#### read in rockstar grp output ####
    stat = read_stat_struc_amiga.read_amiga('out_0.stat')
    
#### plot rockstar output #### Mass in M/h
    logmass = np.log10(sim['Mvir']/h)
    hist, bins = np.histogram(logmass, bins=15)
    normalization = (bins[1:] - bins[:-1])*(box)**3.
    err = 1. + (hist + 0.75)**0.5
    plt.errorbar(10.**((bins[1:]+bins[:-1])/2.), hist/normalization, yerr=err/normalization, linewidth=2, label='Cosmo25V z~10 Rockstar')
    
#### plot rockstar grp output #### Mass in M
    nonzero = stat['m_tot'] > 0.
    logmass = np.log10(stat['m_tot'][nonzero])
    hist, bins = np.histogram(logmass, bins=15)
    normalization = (bins[1:] - bins[:-1])*(box)**3.
    err = 1. + (hist + 0.75)**0.5
    plt.errorbar(10.**((bins[1:]+bins[:-1])/2.), hist/normalization, yerr=err/normalization, linewidth=2, label='Cosmo25V z~10 Stat')
    
    for i in range(len(mffilenames)): 
        f = open(mffilenames[i])
        names = re.sub(r'\([^)]*\)', '', f.readline().strip('\n#  ')).split(' ')
        while '' in names: names.remove('')
        f.close()
        mf = np.genfromtxt(mffilenames[i], names=names, skiprows=1)
        if (mffilenames[i] == 'plk13.z10.behroozi.mf'):
            plt.plot(mf['M']/h, mf['dndlog10m_PLANCK13_z100']*h**3., label=mffilenames[i].split('.')[2], linewidth=2)
        elif (mffilenames[i]=='plk13tweek.z10.behroozi.mf'):
            plt.plot(mf['M']/h, mf['dndlog10m_PLANCK13tweek_z100']*h**3., label='tweeked', linewidth=2)
        else: plt.plot(10.**mf['log_10m']/h, 10.**mf['log_10dn']*h**3., label=mffilenames[i].split('.')[2], linewidth=2)
        
    plt.plot(np.zeros(1e3)+darkmass*1e3, np.linspace(1e-5, 1e4, 1e3), label='1000 DM',  linewidth=2)
    plt.legend()
    plt.yscale('log')
    plt.xscale('log')
    plt.xlabel('M$_{tot}$ [M$_{\odot}$]', fontsize=15)
    plt.ylabel('$\Phi$ [Mpc$^{-3}$ log(M$_{tot}$)]', fontsize=15)
    plt.title('Vulcan Mass Function z~10', fontsize=20)
    plt.xlim(1e5, 2e11) 
    plt.ylim(1e-5, 1e4)
    plt.show()
    pdb.set_trace()
    

if __name__ == '__main__': plot()
