import colorcet as cc
import matplotlib.pyplot as plt
import numpy as np
from plotapp import PlotApp

import scgenerator as sc

params = sc.Parameters(
    wavelength=835e-9,
    width=50e-15,
    peak_power=10e3,
    repetition_rate=20e6,
    shape="sech",
    # fiber
    length=15e-2,
    raman_type="measured",
    beta2_coefficients=[
        -11.830e-27,
        8.1038e-41,
        -9.5205e-56,
        2.0737e-70,
        -5.3943e-85,
        1.3486e-99,
        -2.5495e-114,
        3.0524e-129,
        -1.7140e-144,
    ],
    gamma=0.11,
    # simulation
    tolerated_error=1e-8,
    wavelength_window=[300e-9, 1500e-9],
    t_num=8192,
    quantum_noise=False,
    z_num=128,
)


def compute_manual():
    spec0 = params.compute("spec_0")
    w_c, w0, gamma = params.compute("w_c", "w0", "gamma")
    p = params.compile()
    print(p.dt)
    beta_op = sc.operators.constant_polynomial_dispersion(
        params.beta2_coefficients, w_c, params.compute("dispersion_ind")
    )
    linear = sc.operators.envelope_linear_operator(
        beta_op,
        # sc.operators.constant_quantity(0),
        sc.operators.constant_quantity(0),
    )
    # with PlotApp() as app:
    #     o = np.argsort(p.l)
    #     sax = app["spectrum"]
    #     sax.set_line_data("beta2", p.l[o], beta_op(0)[o].imag)
    # return

    # spm = sc.operators.envelope_spm(0)
    # plt.plot(p.l, sc.abs2(p.spec_0))
    # plt.plot(p.l, sc.abs2(np.fft.fft(spm(np.fft.ifft(p.spec_0), 0))), ls=":")
    # plt.xlim(100e-9, 1500e-9)
    # plt.show()
    # return

    # nonlinear = sc.operators.envelope_nonlinear_operator(
    #     gamma_op=sc.operators.constant_quantity(params.gamma),
    #     ss_op=sc.operators.constant_quantity(w_c / w0),
    #     spm_op=sc.operators.envelope_spm(0),
    #     raman_op=sc.operators.no_op_time(params.t_num),
    # )

    hr_w = params.compute("hr_w")

    def nonlinear(spec, z):
        field = np.fft.ifft(spec)
        field2 = sc.abs2(field)
        fr = 0.18
        return (
            -1j
            * gamma
            * (1 + w_c / w0)
            * np.fft.fft(field * ((1 - fr) * field2 + fr * np.fft.ifft(hr_w * np.fft.fft(field2))))
        )

    above0 = p.l > 0
    linear_arr = np.zeros(p.t_num, dtype=complex)
    w_power_fact = np.array(
        [sc.math.power_fact(w_c[above0], k) for k in range(2, len(p.beta2_coefficients) + 2)]
    )
    for i, wn in reversed(list(enumerate(w_power_fact))):
        print(i, p.beta2_coefficients[i])
        linear_arr[above0] += p.beta2_coefficients[i] * sc.math.power_fact(w_c[above0], i + 2)
    linear_arr *= -1j

    def linear(_):
        return linear_arr

    prop = sc.propagation("examples/dudley_manual.zip", params)
    z = []
    for i, (spec, stat) in enumerate(
        sc.solve43(spec0, linear, nonlinear, params.length, 1e-6, 1e-6, 0.9, h_const=20e-6)
    ):
        if i % 20:
            continue
        print(stat, end="\r")
        prop.append(spec)
        z.append(stat["z"])


def compute_auto():
    sc.compute(params, True, "examples/dudley2006")


def plot():
    prop = sc.propagation("examples/dudley_manual.zip")
    newprop = sc.propagation("examples/dudley2006.zip")
    fig, (top, bot) = plt.subplots(2, 2)
    z = np.linspace(0, prop.parameters.length, len(prop))
    sc.plotting.summary_plot(prop[:], z, axes=top)
    sc.plotting.summary_plot(newprop[:], newprop.parameters.z_targets, axes=bot)
    plt.show()


if __name__ == "__main__":
    # compute_auto()
    plot()