from pyon.Classes.hadron import PseudoscalarMeson
from pyon.Classes import lattice
#from pyon.Error import error_reduce
#import numpy as np
FIT=True
my_lattice = lattice.Lattice24c()
d = '/temp/srd1g10/results/24c/mesons/0.01_0.0613/'
folders = (d+'pickled.g15.g15.p16/', d+'pickled.g7.g7.p16/', d+'pickled.g7.g15.p16/')

mes = PseudoscalarMeson.load_from_folder(folders, fit_range=(8, 31),
                                         name="(0.01, 0.0613)",
                                         lattice=my_lattice,
                                         fit_type='simultaneous',
                                         mass_guess=0.4,
                                         covariant=True,
                                         correlated=False, frozen_error=True,
                                         num_bins=2, verbose=False)

def get_mass():
    mes.fit()
    return mes.get_jackknife_lists()

if __name__== "__main__":
    if FIT:
        mes.fit()
        print mes.get_fit_params()
        print mes.get_fit_params_latex()
    else:
        mes.plot_graph()

import numpy as np
from pyon.IO import prettify
FIT=True
my_lattice = lattice.Lattice24c()
d = '/temp/srd1g10/results/24c/mesons/0.005_0.04/pickled.dispersion/'
folders = [d+'p0/', d+'p1/']
fit_ranges = [(8, 31), (8, 27)]
guesses = [0.329, 0.42]
names = ['p0', 'p1']
jl = []
for folder, fit_range, guess, name in zip(folders, fit_ranges, guesses, names):

    mes = PseudoscalarMeson.load_from_folder(folder, fit_range=fit_range,
                                             name="(0.005, 0.04) " + name,
                                             lattice=my_lattice,
                                             fit_type='individual',
                                             mass_guess=guess,
                                             covariant=True,
                                             correlated=False, frozen_error=True,
                                             num_bins=2, verbose=False)
    #mes.plot_graph()
    mes.fit()
    jl.append(mes.get_jackknife_lists())

jl = np.array(jl)

def get_c():

    jl_2 = jl * jl
    masses = np.average(jl, axis=1) # central masses
    m_2 = masses * masses
    L = 24

"""
Fit a meson with PP only
"""
from pyon.Classes.hadron import PseudoscalarMeson
from pyon.Classes import lattice
from pyon.Error import error_reduce
import numpy as np
import os
import pickle
my_lattice = lattice.Lattice24c()
folder = '/home/srd1g10/charm/data/24c/mesons/0.005_0.005/pickled.g15.g15.p16/'

mes = PseudoscalarMeson.load_from_folder(folder, fit_range=(9, 32),
                                         name="0.005_0.005",
                                         lattice=my_lattice,
                                         fit_type='individual',
                                         covariant=True,
                                         correlated=False, frozen_error=True,
                                         num_bins=2, verbose=False)
mes.plot_graph()
#if mes.fit():
#    print mes.get_fit_params()