import numpy as np
import bagpipes as pipes

# Define Ha equivalent width index (these numbers are just an example)
ha_ew = {}
ha_ew["name"] = "Ha_ew"
ha_ew["type"] = "EW"
ha_ew["feature"] = [6553., 6573.] # The spectral region being measured
ha_ew["continuum"] = [[6533., 6553.], [6573., 6593.]] # The side regions to be compared against


# Define index list variable, the idea here is similar to filt_list
index_list = [ha_ew]


# Set up simple model_components dictionary
nebular = {}
nebular["logU"] = -3.

const = {}
const["age_min"] = 0.
const["age_max"] = 1.
const["metallicity"] = 1.
const["massformed"] = 10.

model_comp = {}
model_comp["redshift"] = 0.001
model_comp["constant"] = const
model_comp["nebular"] = nebular


### Example of making a model_galaxy to demonstrate index calculation
model = pipes.model_galaxy(model_comp, index_list=index_list)

print(model.indices) # EW specified above (absorption is positive)


### Example of loading in index data to a galaxy object

# Define load_indices function to return equivalent width value
def load_ha_ew_value(ID):
    indices = np.zeros((1, 2))
    indices[0, 0] = -15. #index value
    indices[0, 1] = 0.1 #index uncertainty
    return indices

# Example load_data function for photometry
def load_phot_data(ID):
    photometry = np.random.randn(12, 2)
    return photometry


# Set up example filt_list
filt_list = np.loadtxt("filters/list_of_filters.txt", dtype="str")


# Create example galaxy object to demonstrate loading index values
galaxy = pipes.galaxy("test", load_phot_data, filt_list=filt_list,
                      load_indices=load_ha_ew_value, index_list=index_list,
                      spectrum_exists=False)

print(galaxy.indices) # Similar idea to galaxy.photometry


### Example of fitting indices

# Set up simple fit_instructions dictionary
nebular = {}
nebular["logU"] = -3.

const = {}
const["age_min"] = (0., 0.1)
const["age_max"] = 1.
const["metallicity"] = (0.1, 2.5)
const["massformed"] = (0., 13.)

fit_info = {}
fit_info["redshift"] = 0.001
fit_info["constant"] = const
fit_info["nebular"] = nebular


# Set up example fit object
fit = pipes.fit(galaxy, fit_info)


# Fit the photometry + index value
fit.fit(verbose=True)