from __future__ import annotations

import numpy as np
from matplotlib import pyplot as plt

from .distributions import BaseDistribution
from .motions import BaseMotion
from .parameter import Parameter, make_filename


def save_spectrum_data(
        timesignal: np.ndarray,
        spectrum: np.ndarray,
        param: Parameter,
        dist: BaseDistribution,
        motion: BaseMotion,
        echo_strings: list[str]
) -> None:
    filename = make_filename(dist, motion)

    header = param.header(sim=True, spec=True)
    header += '\n' + dist.header()
    header += '\n' + motion.header()
    header += '\nx\t' + '\t'.join(echo_strings)

    np.savetxt(filename + '_timesignal.dat', np.c_[param.spec.t_acq, timesignal], header=header)
    np.savetxt(filename + '_spectrum.dat', np.c_[param.spec.freq, spectrum], header=header)

    fig, ax = plt.subplots()
    lines = ax.plot(param.spec.freq, spectrum)
    ax.set_title(f'{dist}, {motion}')
    ax.legend(lines, echo_strings)
    plt.savefig(filename + '_spectrum.png')

    fig1, ax1 = plt.subplots()
    lines = ax1.plot(param.spec.t_acq, timesignal)
    ax1.set_title(f'{dist}, {motion}')
    ax1.legend(lines, echo_strings)
    plt.savefig(filename + '_timesignal.png')