import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from os import listdir
import re

# Not used galaxy spectra from László Dobos
files = listdir(path='./kompozit_dobos_data/data/')
files = np.array(list(filter(lambda x: re.match(r'.*\.spec\..', x), files)))
files.flatten(), files.size

# Reading the files into arrays
def get_spectra_dobos(files):
    spectra = []
    for file in files:
        # Reading text files with Pandas
        data = pd.read_csv('kompozit_dobos_data/data/%s' % file, sep="\t", header=None)
        # Converting from scientific notation
        data = data.apply(pd.to_numeric, args=('coerce',))
        # Using the wavelength and the flux only
        spectra.append(data.values[1:,0:2])
    return np.array(spectra)

spectra_dobos = get_spectra_dobos(files)

# Plotting all in one
plt.figure(figsize=(15.,9.))
plt.xlabel('Wavelength [Angstrom]')
plt.ylabel('Flux []')
plt.title('Galaxy spectra')
# Filter spectrum
plt.xlim(3000,9500)
for i in range(len(spectra_dobos)):
    plt.plot(spectra_dobos[i][:,0][::100], spectra_dobos[i][:,1][::100])

def get_lephare_spectra():
    # Reading the Le Phare sepctra
    dirs = listdir(path='lephare_dev/sed/GAL/')
    files = np.array([])
    spectra = []
    filenames = []
    for dr in dirs:
        # Some kind of bug, dirs read with ._ at the beginning
        if(dr[0] != "."):
            files = listdir(path='lephare_dev/sed/GAL/%s/' % dr)
            file_names = np.array(list(filter(lambda x: re.match(r'.*\.sed', x), files)))
            for file in file_names:
                # Some kind of bug, dirs read with ._ at the beginning
                if(file[0] != "."):
                    # Due to erros in file formatting
                    try:
                        # Reading in %.3f format
                        pd.set_option('display.float_format', lambda x: '%.3f' % x)
                        # Read with any whitespace separator (still some errors)
                        data = pd.read_csv('lephare_dev/sed/GAL/%s/%s' % (dr, file), sep=r"\s*",
                                           header=None, engine='python')
                        # Numeric 
                        data = data.apply(pd.to_numeric, args=('coerce',))
                        data = data.dropna(axis='rows')
                        if(len(data.values.shape) > 1 and data.values.shape[0] != 0 and data.values.shape[1] != 0):
                            filenames.append('lephare_dev/sed/GAL/%s/%s' % (dr, file))
                            spectra.append(data.values[:, 0:2])
                    finally:
                        # Progress bar
                        print('=', end='')
    # End progress bar
    print(">\t", len(spectra))
    return np.array(spectra), np.array(filenames)

lephare_spectra, filenames = get_lephare_spectra()

i = 1
plt.figure(figsize=(24, 32))
for j in range(4, lephara_spectra.size, 33):
    plt.subplot(int('43%d' % i))
    i += 1
    plt.xlim(2900,13000)
    plt.title('Data from: %s' % filenames[j])
    plt.xlabel('Wavelength [Angstrom]')
    plt.ylabel('Flux []')
    plt.plot(lephare_spectra[j][:,0], lephare_spectra[j][:,1])
plt.show()

colors = ['u', 'g', 'r', 'i', 'z']
filters = np.array([pd.read_csv('./SDSS_%s\'.txt' % c, sep=r"\s", header=None, engine='python').values for c in colors])

plt.figure(figsize=(16.,18.))
for ind, color in enumerate(colors):
    plt.subplot(int("32%d" % (ind+1)))
    plt.title('Filter for color %s' % color)
    plt.xlabel("Wavelength [Angstrom]")
    plt.ylabel("Flux []")
    plt.plot(filters[ind][:,0], filters[ind][:,1])
plt.show()

def get_spectrum(spectrum, color_filter):
    # min/max wavelength
    min_lambda = color_filter[:,0][0]
    max_lambda = color_filter[:,0][-1]
    
    # find nearest min/max indices
    min_nearest_ind = np.abs(spectrum[:,0]-min_lambda).argmin()
    max_nearest_ind = np.abs(spectrum[:,0]-max_lambda).argmin()
    
    # integral of response function
    c = 3*10**8
    bin_width = color_filter[:,0][1] - color_filter[:,0][0]
    filter_integral = c*np.sum(color_filter[:,1]/color_filter[:,0])*10**10 # dimension correction
    
    flux = 0.
    for ind, val in enumerate(spectrum[:,0][min_nearest_ind:max_nearest_ind]):
        nearest_ind = np.abs(color_filter[:,0] - val).argmin()
        flux += color_filter[:, 1][nearest_ind]*spectrum[:,1][ind]*val*10**(-10) # dimension correction
    return -2.5*np.log10(flux/filter_integral) - 48.60

def build_dataframe_from_template(_spectra, _colors, _N=50):
    data = []
    spectra = _spectra.copy()
    colors = np.append(_colors, ['redshift'])
    for z in np.linspace(0.,1.,_N):
        for j in range(spectra.size):
            row = []
            spectrum = spectra[j].copy()
            spectrum[:,1] *= (1.+z)
            for i in range(filters.size):
                # magnitudes and redshift
                row.append(get_spectrum(spectrum, filters[i]))
            row.append(z)
            data.append(row)
    # Building dataframe \w column names
    return pd.DataFrame(data, columns=colors)

template_data = build_dataframe_from_template(lephara_spectra, colors, 10)

template_data

template_data = pd.read_csv('template_dataset_150_steps_from0to1_approx45000lines.csv')
template_data.head(n=3)

measured_data = pd.read_csv('SDSSDR7-5-percent-err-limit-250000-line.csv')
measured_data['redshift'] = measured_data['z'] # rename column z to redshift (name mismatch)
measured_data['r'] = measured_data['m_r']
measured_data.head(n=3)

template_data['ug'] = template_data['u'] - template_data['g']
template_data['gr'] = template_data['g'] - template_data['r']
template_data['ri'] = template_data['r'] - template_data['i']
template_data['iz'] = template_data['i'] - template_data['z']
template_data.head(n=3)

def find_nearest_template_in_measured_data(template_data, measured_data, indices):
    for row in measured_data[indices].values:
        print('MEASURED DATA:\t', row)
        min_ind = np.argmin([np.sqrt(np.sum(np.square(row - val))) for val in template_data[indices].values])
        print('BEST TEMPLATE DATA: ',template_data[indices].values[min_ind, :], ' - MIN INDEX :', min_ind)
        break

find_nearest_template_in_measured_data(template_data, measured_data, ['ug', 'gr', 'ri', 'iz', 'redshift' ])

def find_nearest_template_in_measured_data(template_data, measured_data, indices, batch=100):
    splitted_datasets = np.split(measured_data[indices].values, batch)
    dataframes = []
    for num, dataset in enumerate(splitted_datasets):
        pd.set_eng_float_format(accuracy=5, use_eng_prefix=True)
        best_template_data = []
        for row in dataset:
            min_ind = np.argmin(np.linalg.norm(template_data[indices].values - row, axis=1))
            best_template_data.append(template_data[indices].values[min_ind, :])
        dataframe = pd.DataFrame(data=np.array(best_template_data), columns=indices)
        dataframe.to_csv('dataframes/dataframe%d.csv' % num)
    return

find_nearest_template_in_measured_data(template_data, measured_data, ['r', 'ug', 'gr', 'ri', 'iz', 'redshift' ])