import numpy as np
import Apodization
import pandas as pd
import os
from scipy.interpolate import interp1d

'''dft_2DMAP computes the time resolved emission map from the interferometric map
    INPUTS:
    - map: interferometric map each row must be the histogram correspondent to a specific delay as specified by pos
    - pos: position axis
    - start_wave: start wavelength from which to calculate the spectra
    - end_wave: end wavelength until which to calculate the spectra
    - samples: number of spectral samples to calculate
    - apodization width: normalized standard deviation with respect to the
    - position axis length of the apodization function
    OUTPUTS:
    - MAP: time resolved emission map
    - wave: wavelength axis calibrated
'''


def dft_2DMAP(map, pos, start_wave, end_wave, samples, apodization_width):
    # Load spectral calibration file
    items = os.listdir(".")

    for names in items:
        if names.endswith("parameters_cal.txt"):
            filename = names

    ref = pd.read_csv(filename, sep="\t", header=None)
    first_row = (ref.iloc[0])
    second_row = (ref.iloc[1])
    wavelength = first_row.to_numpy(dtype='float64')
    reciprocal = second_row.to_numpy(dtype='float64')

    # Compute the frequency limits from the calibration file
    fn = interp1d(1 / wavelength, reciprocal, kind="linear")
    start_freq = fn(1 / end_wave)
    end_freq = fn(1 / start_wave)

    # Create differential axis
    dpos1 = (np.diff(pos))
    dpos = np.append(dpos1, 0)

    # Create frequency axis
    freq = np.linspace(start_freq, end_freq, samples)

    # Perform the DFT with the apodization
    size = len(map[1])
    a = pos.reshape(-1, 1)
    exp_map = np.exp(-1j * 2 * np.pi * a * freq)
    MAP = np.zeros((samples, size))

    for k in range(size):
        interf, window = Apodization.apodization(map[:, k] - map[:, k].mean(), pos, apodization_width)
        a = dpos * interf
        b = a.reshape(1, -1)
        MAP[:, k] = abs(b.dot(exp_map))

    # Compute the calibrated wavelength axis
    fn = interp1d(reciprocal, 1 / wavelength, kind="linear")
    wave = 1 / fn(freq)

    return MAP, wave
