"""
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.
"""

import sys
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):

    # 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_sigma = 0.15  # std for gauss broadening
        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(1.5, 5)
        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 += ".pdf"  # change here for pdf or png output
        output_path = output_folder / output_name
        plt.savefig(output_path)
        plt.close(fig)  # close plot


# getting commandline input and pass it to function
if __name__ == "__main__":
    a = str(sys.argv[1])
    b = str(sys.argv[2])
    exspectrum_plotter(a, b)

# python exspectrum_plotter.py "C:\Path\to\Calculations" "C:\path\to\results"