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 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__":
    main()