import numpy as np
import prospect
from prospect.sources import CSPSpecBasis
from prospect.models import priors, sedmodel, SpecModel, PolySpecModel
from astropy.io import fits
from astropy.table import Table
import matplotlib.pyplot as plt
from prospect.models.templates import TemplateLibrary
import os

### READ IN SPECTRAL DATA

def read(filename):
    
    hdu_list = fits.open(filename, memmap = True)
    data = Table(hdu_list[1].data)
    
    return data

spec = read('fits_files/NSC_VCC1261_R27_smooth.fits')

### BUILD OBS

def build_obs(snr=10, ldist=10.0, **extras):
    
    from prospect.utils.obsutils import fix_obs
    
    # The obs dictionary, empty for now
    obs = {}
    
    obs["filters"] = None #sedpy.observate.load_filters(filternames)
    
    obs["wavelength"] = spec['WAVELENGTH']
    obs["spectrum"] = spec['FLUX']
    obs['unc'] = spec['ERR']
    obs['mask'] = spec['MASK']
    
    obs = fix_obs(obs)
    
    return obs

run_params = {}
run_params["snr"] = 10.0
run_params["ldist"] = 10.0

### BUILD SPS

def build_sps(zcontinuous=1, **extras):
    """
    :param zcontinuous: 
        A vlue of 1 insures that we use interpolation between SSPs to 
        have a continuous metallicity parameter (`logzsol`)
        See python-FSPS documentation for details
    """
    from prospect.sources import CSPSpecBasis
    sps = CSPSpecBasis(zcontinuous=zcontinuous)
    return sps

run_params["zcontinuous"] = 1

### BUILD MODEL

def build_model(object_redshift=None, ldist=10.0, fixed_metallicity=None, add_duste=False, **extras):
    
    model_params = TemplateLibrary['ssp']
    model_params.update(TemplateLibrary['optimize_speccal'])
    
    # model_params["lumdist"] = {"N": 1, "isfree": False, "init": ldist, "units":"Mpc"}
    
    model_params['zred']['isfree'] = True
    model_params['dust2']['isfree'] = False
    
    #optimize_speccal mods
    model_params['spec_norm']['isfree'] = False
    
    #inits for fixed parameters
    model_params["dust2"]["init"] = 0.0
    model_params["zred"]["init"] = 0.00625
    
    #inits for free parameters
    model_params["mass"]["init"] = 1e5
    model_params["logzsol"]["init"] = -0.5
    model_params["tage"]["init"] = 5.0
    
    #priors
    model_params["mass"]["prior"] = priors.LogUniform(mini=1e0, maxi=1e10)
    model_params["logzsol"]["prior"] = priors.TopHat(mini=-1.0, maxi=1.0)
    model_params["tage"]["prior"] = priors.TopHat(mini=1.0, maxi=13.8)
    
    #minimum vals for free parameters
    model_params["mass"]["disp_floor"] = 1e0
    model_params["tage"]["disp_floor"] = 1.0
    
    model = sedmodel.PolySpecModel(model_params)
    
    # print(model.theta)
    print(model_params.keys())
    
    return model

### BUILD EVERYTHING

def build_all(**run_params):
    
    return (build_obs(**run_params), build_sps(**run_params), build_model(**run_params))

obs, sps, model = build_all(**run_params)

### GENERATE INITIAL MODEL

theta = model.theta.copy()
initial_spec, initial_phot, initial_mfrac = model.predict(theta, obs=obs, sps=sps)

theta

### LIKELIHOOD FUNCTION FOR MODEL FITTING

from prospect.likelihood import lnlike_spec, lnlike_phot, write_log

verbose = False
def lnprobfn(theta, model=None, obs=None, sps=None, 
             nested=False, verbose=verbose):
    """
    Given a parameter vector, a model, a dictionary of observational 
    data, and an sps object, return the ln of the posterior. 
    """

    # Calculate prior probability and exit if not within prior
    # Also if doing nested sampling, do not include the basic priors, 
    # since the drawing method includes the prior probability
    lnp_prior = model.prior_product(theta, nested=nested)
    if not np.isfinite(lnp_prior):
        return -np.infty
        
    # Generate "mean" model
    spec, phot, mfrac = model.mean_model(theta, obs, sps=sps)
 
    # Calculate likelihoods
    lnp_spec = lnlike_spec(spec, obs=obs)
    lnp_phot = lnlike_phot(phot, obs=obs)

    return lnp_prior + lnp_phot + lnp_spec

run_params["verbose"] = verbose

from prospect.likelihood import chi_spec, chi_phot
def chivecfn(theta):
    """A version of lnprobfn that returns the simple uncertainty 
    normalized residual instead of the log-posterior, for use with 
    least-squares optimization methods like Levenburg-Marquardt.
    
    It's important to note that the returned chi vector does not 
    include the prior probability.
    """
    lnp_prior = model.prior_product(theta)
    if not np.isfinite(lnp_prior):
        return np.zeros(model.ndim) - np.infty

    # Generate mean model
    try:
        spec, phot, x = model.mean_model(theta, obs, sps=sps)
    except(ValueError):
        return np.zeros(model.ndim) - np.infty

    chispec = chi_spec(spec, obs)
    chiphot = chi_phot(phot, obs)
    return np.concatenate([chispec, chiphot])

from prospect.fitting import lnprobfn
from prospect.fitting import fit_model

### MINIMIZATION

# --- start minimization ----
run_params["dynesty"] = False
run_params["emcee"] = False
run_params["optimize"] = True
run_params["min_method"] = 'lm'
# We'll start minimization from "nmin" separate places, 
# the first based on the current values of each parameter and the 
# rest drawn from the prior.  Starting from these extra draws 
# can guard against local minima, or problems caused by 
# starting at the edge of a prior (e.g. dust2=0.0)
run_params["nmin"] = 2

output = fit_model(obs, model, sps, lnprobfn=lnprobfn, **run_params)

print("Done optmization in {}s".format(output["optimization"][1]))

(results, topt) = output["optimization"]
# Find which of the minimizations gave the best result, 
# and use the parameter vector for that minimization
ind_best = np.argmin([r.cost for r in results])
print(ind_best)
theta_best = results[ind_best].x.copy()
print(theta_best)

### EMCEE SAMPLING

# Set this to False if you don't want to do another optimization
# before emcee sampling (but note that the "optimization" entry 
# in the output dictionary will be (None, 0.) in this case)
# If set to true then another round of optmization will be performed 
# before sampling begins and the "optmization" entry of the output
# will be populated.
run_params["optimize"] = False
run_params["emcee"] = True
run_params["dynesty"] = False
# Number of emcee walkers
run_params["nwalkers"] = 128
# Number of iterations of the MCMC sampling
run_params["niter"] = 512
# Number of iterations in each round of burn-in
# After each round, the walkers are reinitialized based on the 
# locations of the highest probablity half of the walkers.
run_params["nburn"] = [32, 64, 128, 256]

output = fit_model(obs, model, sps, lnprobfn=lnprobfn, **run_params)
print('done emcee in {0}s'.format(output["sampling"][1]))

### SAVE RESULTS OF EMCEE TO H5 FILE

from prospect.io import write_results as writer
hfile = "Emcee_h5_files/demo_emcee_mcmc.h5"
writer.write_hdf5(hfile, run_params, model, obs,
                  output["sampling"][0], output["optimization"][0],
                  tsample=output["sampling"][1],
                  toptimize=output["optimization"][1])

print('Finished')

from prospect import prospect_args
# - Parser with default arguments -
parser = prospect_args.get_parser()
# - Add custom arguments -
parser.add_argument('--add_duste', action="store_true",
                    help="If set, add dust emission to the model.")
parser.add_argument('--ldist', type=float, default=10,
                    help=("Luminosity distance in Mpc. Defaults to 10"
                          "(for case of absolute mags)"))
args, _ = parser.parse_known_args()
cli_run_params = vars(args)
print(cli_run_params)

### DISPLAY EMCEE RESULTS

import prospect.io.read_results as reader
results_type = "emcee"
# grab results (dictionary), the obs dictionary, and our corresponding models
# When using parameter files set `dangerous=True`
result, obs, _ = reader.results_from("Emcee_h5_files/demo_emcee_mcmc.h5".format(results_type), dangerous=True)

#The following commented lines reconstruct the model and sps object, 
# if a parameter file continaing the `build_*` methods was saved along with the results
#model = reader.get_model(result)
#sps = reader.get_sps(result)

# let's look at what's stored in the `result` dictionary
print(result.keys())

### WALKER CHAINS

if results_type == "emcee":
    chosen = np.random.choice(result["run_params"]["nwalkers"], size=10, replace=False)
    # tracefig = reader.traceplot(result, figsize=(20,10), chains=chosen)
    tracefig = reader.traceplot(result, figsize=(20,10), chains=np.arange(0, result["run_params"]["nwalkers"]), c = 'k')
else:
    tracefig = reader.traceplot(result, figsize=(20,10), c = 'k')

### CORNER PLOT

# maximum a posteriori (of the locations visited by the MCMC sampler)
imax = np.argmax(result['lnprobability'])
if results_type == "emcee":
    i, j = np.unravel_index(imax, result['lnprobability'].shape)
    theta_max = result['chain'][i, j, :].copy()
    thin = 5
else:
    theta_max = result["chain"][imax, :]
    thin = 1

# print('Optimization value: {}'.format(theta_best))
print('MAP value: {}'.format(theta_max))
print(result['chain'].shape)
cornerfig = reader.subcorner(result, start=0, thin=thin, truths=theta_max, 
                             fig=plt.subplots(4,4,figsize=(8,8))[0])