"""
Turbomole-Spectrum-Plotter
(c) 2022 Lukas Schank

This script will run through all subfolders of the given folder and plot excitation spectra of TDDFT calculations.

Call on command line like this:
python exspectrum_plotter.py --data_folder "C:\Path\to\Calculations" --output_folder "C:\path\to\results" --calc_sigma 0.15 --output_type "pdf" --save_csv False

calc_sigma output_type and save_csv are optional.
"""

import argparse
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path


def exspectrum_plotter(data_folder: str, output_folder: str, calc_sigma=0.15, output_type="pdf", save_csv=False):

    """
    Iterates through all subfolders of data_folder and plots data of exspectrum files.
    :param data_folder: path to input folder
    :param output_folder: path to output folder, gets created if necessary
    :param calc_sigma: standard deviation of gauss broadening for calculated transitions
    :param output_type: graphs can be saved as "pdf" or "png"
    :param save_csv: write interpolated spectrum as csv
    :return: files to the folder specified in output_folder
    """

    # convert input to path objects
    data_folder = Path(data_folder)
    output_folder = Path(output_folder)
    if not output_folder.exists():
        output_folder.mkdir()

    # convert data to plot
    for data_file in data_folder.glob('**/exspectrum'):  # path.glob iterates through given folder for pattern
        # read data
        rawdata = pd.read_csv(
            data_file,
            sep='\s+',  # seperator is one or more spaces
            skiprows=2,  # skip header rows
            header=None  # prevent getting first numbers as header
        )
        data_energy = np.array(rawdata.iloc[:, 3])
        data_osc_strength = np.array(rawdata.iloc[:, 6])

        # calculate gauss broadening for spectrum
        def gauss_spectrum(energy, osc_strength, sigma, gauss_energy_range):
            gauss_osc_strength = []
            for E_i in gauss_energy_range:
                tot = 0
                for E_j, osc in zip(energy, osc_strength):
                    tot += osc * np.exp(-(((E_j - E_i) / sigma) ** 2))
                gauss_osc_strength.append(tot)
            return gauss_osc_strength

        calc_energy_range = np.linspace(0, 6, num=500, endpoint=True)  # x values for calculated spectrum
        calc_osc_strength = gauss_spectrum(data_energy, data_osc_strength, calc_sigma, calc_energy_range)

        # create plot
        fig, ax = plt.subplots(dpi=300, figsize=(6, 4))
        ax.plot(calc_energy_range, calc_osc_strength, "-k")  # plot calculated spectrum
        for plt_energy, plt_osc_strength in zip(data_energy, data_osc_strength):
            ax.plot((plt_energy, plt_energy), (0, plt_osc_strength), c="k")  # plot lines from data
        ax.set_xlabel("Energy / eV", fontsize=16)
        ax.set_ylabel("Osc. Strength / a.u.", fontsize=16)
        ax.xaxis.set_tick_params(labelsize=14, width=1.5)
        ax.yaxis.set_tick_params(labelsize=14, width=1.5)
        for axis in ['top', 'bottom', 'left', 'right']:
            ax.spines[axis].set_linewidth(1.5)
        ax.set_xlim(0, 6)
        plt.tight_layout()

        # save plot
        data_name = Path(data_file)
        output_name = data_name.parts[data_folder.parts.__len__()]
        for name in data_name.parts[(data_folder.parts.__len__() + 1):]:
            output_name += "_" + name
        output_name_plot = output_name + "." + output_type
        output_path = output_folder / output_name_plot
        plt.savefig(output_path)
        plt.close(fig)  # close plot

        # save calculated data
        if save_csv:
            csv_output = pd.DataFrame({'Energy': calc_energy_range.tolist(), 'Osc. Strength': calc_osc_strength})
            output_name_csv = output_name + ".csv"
            output_path_2 = output_folder / output_name_csv
            csv_output.to_csv(output_path_2)


# getting commandline input and pass it to function
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_folder', type=str, required=True)
    parser.add_argument('--output_folder', type=str, required=True)
    parser.add_argument('--calc_sigma', type=float, required=False)
    parser.add_argument('--output_type', type=str, required=False)
    parser.add_argument('--save_csv', type=bool, required=False)
    args = parser.parse_args()
    exspectrum_plotter(**vars(args))