import numpy as np


def feng_aif(t, tau, A1, A2, A3, L1, L2, L3):
    cp = (
        (A1 * (t - tau) - A2 - A3) * np.exp(-L1 * (t - tau))
        + A2 * np.exp(-L2 * (t - tau))
        + A3 * np.exp(-L3 * (t - tau))
    )
    cp[: np.sum(t < tau)] = 0
    return cp


def frame_average(y, t_s, frame_durations_s):
    edges = np.concatenate([[0], np.cumsum(frame_durations_s)])
    out = np.empty(len(frame_durations_s))
    for i, (a, b) in enumerate(zip(edges[:-1], edges[1:])):
        m = (t_s >= a) & (t_s < b)
        if np.any(m):
            out[i] = np.trapz(y[m], t_s[m]) / (b - a)
        else:
            out[i] = np.interp(0.5 * (a + b), t_s, y)
    return out


def generate_aif(
    t_fine_min,
    t_fine_s,
    frame_durations_s,
    dt_s,
    *,
    Tinj_s,
    tau,
    A1,
    r2,
    r3,
    L1,
    L2,
    L3,
):
    k = max(1, int(round(Tinj_s / dt_s)))
    kernel = np.ones(k) / k
    A2, A3 = r2 * A1, r3 * A1

    cp = feng_aif(t_fine_min, tau, A1, A2, A3, L1, L2, L3)
    cp_inj = np.convolve(cp, kernel, mode="same")
    return frame_average(cp_inj, t_fine_s, frame_durations_s)


def generate_time(*, frame_durations_s, dt_s):
    durs = []
    for n_frames, frame_dur in frame_durations_s:
        durs.append(np.full(n_frames, frame_dur))

    frame_durations_s = np.concatenate(durs)
    edges_s = np.concatenate([[0], np.cumsum(frame_durations_s)])
    mids_s = 0.5 * (edges_s[:-1] + edges_s[1:])

    t_fine_s = np.arange(0, edges_s[-1] + dt_s, dt_s)
    t_fine_min = t_fine_s / 60
    time_min = mids_s / 60

    return t_fine_min, t_fine_s, time_min, frame_durations_s, dt_s


def gaussian_noise(signal, rng, *, factor):
    amplitude_noise = rng.normal(loc=0, scale=signal)
    return signal + factor * amplitude_noise
