import multiprocessing
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

import scgenerator as sc

PARAMS = Path("./examples/noisy.toml")
SEED = 2564


def path(i) -> Path:
    return Path(f"build/{Path(__file__).stem}_{i}.zip")


def propagate(i):
    np.random.seed(SEED + i)
    params = sc.Parameters.load("./examples/noisy.toml")

    sc.propagation(path(i), params.compile()).simulate()


def propagate_all(n):
    to_do = [i for i in range(n) if not path(i).exists()]
    with multiprocessing.Pool(4) as pool:
        pool.map(propagate, to_do)

    spec, props = sc.propagation_series([path(i) for i in range(n)])

    return spec, props


def quick_test():
    t = sc.tspace(dt=1e-15, t_num=2048)
    spec_0 = np.fft.fft(10e3 * np.exp(-((t / 70e-15) ** 2)))
    spec = sc.Spectrum(
        [np.exp(2j * np.pi * np.random.rand()) + spec_0 for _ in range(20)],
        sc.wspace(t) + sc.units.nm_rads(800),
        t,
    )
    _, (top, bot) = plt.subplots(2, 1, constrained_layout=True, height_ratios=[1, 5], sharex=True)
    bot.plot(spec.wl_disp * 1e9, spec[0].wl_int)
    top.plot(spec.wl_disp * 1e9, spec.coherence())
    top.set_xlim(750, 850)
    bot.set_yscale("log")
    plt.show()


def main():
    n = 1
    spec, props = propagate_all(n)
    print(spec.shape)

    wl, ind, _ = sc.PlotRange(
        spec.wl_disp[spec.wl_disp > 0].min() * 1e9,
        props.parameters.wavelength_window[1] * 1e9,
        "nm",
    ).sort_axis(spec.wl_disp)

    for i in range(spec.shape[1]):
        fig, (left, right) = plt.subplots(1, 2)
        for s in spec[:, i].time_int:
            left.plot(spec.t, s)
        for s in spec[:, i].wl_int:
            right.plot(wl, s[ind])
        plt.show()


if __name__ == "__main__":
    quick_test()
    # main()