# -*- coding: utf-8 -*-
"""
A subroutine functions to create different excitation and refocusing
pulses accompanied by combined SS and spoil gradients.
Requires the params structure as input.

@author: petrm
"""


# import
from pypulseq.make_sinc_pulse import make_sinc_pulse
from pypulseq.make_trapezoid import make_trapezoid
from pypulseq.make_extended_trapezoid import make_extended_trapezoid
from pypulseq.calc_duration import calc_duration
import numpy as np


# def tse_excitation_grad(params, scanner_parameters, area_gz_spoil, flip180):
#
#    return

def refocusing_grad(params, scanner_parameters, gz90_area, flip180, rf180_phase, t_refwd, spoil_duration, united: bool):
    # Create 180 degree SS refocusing pulse with SS and spoiled gradients
    rf180, gz_ref, _ = make_sinc_pulse(
        flip_angle=flip180,
        system=scanner_parameters,
        duration=params.t_ref,
        slice_thickness=params.sl_thkn,
        apodization=params.apodization,
        time_bw_product=round(params.t_BW_product_ref, 8),
        phase_offset=rf180_phase,
        use="refocusing",
        return_gz=True, 
        delay = scanner_parameters.rf_dead_time
    )


    gz180 = make_trapezoid(channel="z", system=scanner_parameters, amplitude=gz_ref.amplitude,
                           flat_time=t_refwd, rise_time=params.dG)

    if spoil_duration == 'min':
        gz_spoil_dur_min = params.dG + (gz90_area)/scanner_parameters.max_grad
        gz_spoil_dur_min = np.ceil(gz_spoil_dur_min / scanner_parameters.grad_raster_time) * scanner_parameters.grad_raster_time

        if gz_spoil_dur_min < 2 * params.dG:
            gz_spoil_dur_min = 2 * params.dG + scanner_parameters.grad_raster_time

        gz_spoil1 = make_trapezoid(channel='z', system=scanner_parameters, area= gz90_area,
                                   duration=gz_spoil_dur_min, rise_time=params.dG)
        gz_spoil2 = make_trapezoid(channel='z', system=scanner_parameters, area= gz90_area,
                                   duration=gz_spoil_dur_min, rise_time=params.dG)

    else:
        spoil_duration = float(spoil_duration)
        spoil_duration = np.ceil(spoil_duration / scanner_parameters.grad_raster_time) * scanner_parameters.grad_raster_time

        gz_spoil1 = make_trapezoid(channel='z', system=scanner_parameters, area= gz90_area,
                                   duration=spoil_duration, rise_time=params.dG)
        gz_spoil2 = make_trapezoid(channel='z', system=scanner_parameters, area= gz90_area,
                                   duration=spoil_duration, rise_time=params.dG)

    # SS refocusing gradient with spoilers

    gz_sp1_times = np.array(
        [
            0,
            gz_spoil1.rise_time,
            gz_spoil1.rise_time + gz_spoil1.flat_time,
            gz_spoil1.rise_time + gz_spoil1.flat_time + gz_spoil1.fall_time
        ]
    )
    gz_sp1_amp = np.array(
        [
            0,
            gz_spoil1.amplitude,
            gz_spoil1.amplitude,
            gz180.amplitude
        ]
    )
    gz_sp1 = make_extended_trapezoid(channel='z', system=scanner_parameters, times=gz_sp1_times, amplitudes=gz_sp1_amp)

    gz_sp2_times = np.array(
        [
            0,
            gz180.flat_time
        ]
    )
    gz_sp2_amp = np.array(
        [
            gz180.amplitude,
            gz180.amplitude
        ]
    )
    gz_sp2 = make_extended_trapezoid(channel='z', system=scanner_parameters, times=gz_sp2_times, amplitudes=gz_sp2_amp)

    gz_sp3_times = np.array(
        [
            0,
            gz_spoil2.rise_time,
            gz_spoil2.rise_time + gz_spoil2.flat_time,
            gz_spoil2.rise_time + gz_spoil2.flat_time + gz_spoil2.fall_time
        ]
    )

    gz_sp3_amp = np.array(
        [
            gz180.amplitude,
            gz_spoil2.amplitude,
            gz_spoil2.amplitude,
            0
        ]
    )
    gz_sp3 = make_extended_trapezoid(channel='z', system=scanner_parameters, times=gz_sp3_times, amplitudes=gz_sp3_amp)


    if united:
        return rf180, gz_sp1, gz_sp2, gz_sp3
    else:
        return rf180, gz_spoil1, gz180, gz_spoil2

'''
def readout_grad_x(params, scanner_parameters, spoil_duration, rewind_duration, united: bool, phi_radial):
    """
    merged with spoiler gradients readout block in x direction
    phi_radial - angle for radial aquisition
                 (0 - for cartesian)
    @author: petrm
    """

    # Create readout gradient
    readout_time = round(1 / params.BW_pixel, 8)
    k_read = np.double(params.Nf) / np.double(params.FoV_f)
    t_gx = np.ceil(readout_time / scanner_parameters.grad_raster_time) * scanner_parameters.grad_raster_time
    gx = make_trapezoid(channel='x', system=scanner_parameters, flat_area=k_read * np.cos(phi_radial),
                        flat_time=t_gx + 2 * scanner_parameters.adc_dead_time, rise_time=params.dG)

    # generate gx spoiler gradient - G_crr
    if spoil_duration == 'min':
        gx_spoil = make_trapezoid(channel='x', system=scanner_parameters, area=gx.area,
                                  rise_time=params.dG)
        if gx_spoil.flat_time == 0:
            gx_spoil = make_trapezoid(channel='x', system=scanner_parameters, area=gx.area,
                                      rise_time=params.dG, flat_time=params.dG)

    else:
        spoil_duration = float(spoil_duration)
        spoil_duration = (np.ceil(spoil_duration / scanner_parameters.grad_raster_time)
                          * scanner_parameters.grad_raster_time)
        gx_spoil = make_trapezoid(channel='x', system=scanner_parameters, area=gx.area,
                                  duration=spoil_duration, rise_time=params.dG)

    # readout gradient with spoilers

    gx_sp1_times = np.array(
        [
            0,
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil.rise_time) / scanner_parameters.grad_raster_time),
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil.rise_time + gx_spoil.flat_time) / scanner_parameters.grad_raster_time),
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil.rise_time + gx_spoil.flat_time + gx_spoil.fall_time) / scanner_parameters.grad_raster_time)
        ]
    )
    gx_sp1_amp = np.array(
        [
            0,
            gx_spoil.amplitude,
            gx_spoil.amplitude,
            gx.amplitude
        ]
    )
    gx_sp1 = make_extended_trapezoid(channel='x', system=scanner_parameters, times=gx_sp1_times, amplitudes=gx_sp1_amp)

    gx_sp2_times = np.array(
        [
            0,
            scanner_parameters.grad_raster_time * np.round(
                (gx.flat_time) / scanner_parameters.grad_raster_time)
        ]
    )
    gx_sp2_amp = np.array(
        [
            gx.amplitude,
            gx.amplitude
        ]
    )
    gx_sp2 = make_extended_trapezoid(channel='x', system=scanner_parameters, times=gx_sp2_times, amplitudes=gx_sp2_amp)

    gx_sp3_times = np.array(
        [
            0,
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil.rise_time) / scanner_parameters.grad_raster_time),
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil.rise_time + gx_spoil.flat_time) / scanner_parameters.grad_raster_time),
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil.rise_time + gx_spoil.flat_time + gx_spoil.fall_time) / scanner_parameters.grad_raster_time)
        ]
    )

    gx_sp3_amp = np.array(
        [
            gx.amplitude,
            gx_spoil.amplitude,
            gx_spoil.amplitude,
            0
        ]
    )
    gx_sp3 = make_extended_trapezoid(channel='x', system=scanner_parameters, times=gx_sp3_times, amplitudes=gx_sp3_amp)

    if rewind_duration == 'min':
        gx_pre = make_trapezoid(channel="x", system=scanner_parameters, area=gx.area * 1.5,
                                rise_time=params.dG)
        if gx_pre.flat_time == 0:
            gx_pre = make_trapezoid(channel='x', system=scanner_parameters, area=gx.area * 1.5,
                                      rise_time=params.dG, flat_time=params.dG)

    else:
        rewind_duration = float(rewind_duration)
        rewind_duration = (np.ceil(rewind_duration / scanner_parameters.grad_raster_time)
                          * scanner_parameters.grad_raster_time)
        gx_pre = make_trapezoid(channel="x", system=scanner_parameters, area=gx.area * 1.5,
                                rise_time=params.dG, duration=rewind_duration)

    if united:
        return gx_sp1, gx_sp2, gx_sp3, t_gx, gx_pre
    else:
        return gx_spoil, gx, gx_spoil, t_gx, gx_pre
'''

def readout_grad_y(params, scanner_parameters, spoil_duration, rewind_duration,
                   united: bool, phi_radial, phase_encoding_area):
    """
    merged with spoiler gradients readout block in y direction
    phi_radial - angle for radial acquisition
                 (0 - for cartesian)
    phase_encoding_area - additional gradient of phase encoding for spoiler
    @author: petrm
    """

    phase_encoding_area = phase_encoding_area * np.cos(phi_radial)
    # Create readout gradient
    readout_time = round(1 / params.BW_pixel, 8)
    k_read = np.double(params.Nf) / np.double(params.FoV_f)
    t_gx = np.ceil(readout_time / scanner_parameters.grad_raster_time) * scanner_parameters.grad_raster_time
    if phi_radial == 0.0:
        gx = make_trapezoid(channel='y', system=scanner_parameters, area=0,
                            flat_time=t_gx + 2 * scanner_parameters.adc_dead_time, rise_time=params.dG)
    else:
        gx = make_trapezoid(channel='y', system=scanner_parameters, flat_area=k_read * np.sin(phi_radial),
                            flat_time=t_gx + 2 * scanner_parameters.adc_dead_time, rise_time=params.dG)

    # generate gx spoiler gradient - G_crr
    if spoil_duration == 'min':
        gx_spoil_dur_min = params.dG + (gx.area + abs(phase_encoding_area))/scanner_parameters.max_grad
        gx_spoil_dur_min = np.ceil(gx_spoil_dur_min / scanner_parameters.grad_raster_time) * scanner_parameters.grad_raster_time
        if gx_spoil_dur_min < 2 * params.dG:
            gx_spoil_dur_min = 2 * params.dG + scanner_parameters.grad_raster_time

        gx_spoil1 = make_trapezoid(channel='y', system=scanner_parameters, area=gx.area + phase_encoding_area,
                                   duration=gx_spoil_dur_min, rise_time=params.dG)
        gx_spoil2 = make_trapezoid(channel='y', system=scanner_parameters, area=gx.area - phase_encoding_area,
                                   duration=gx_spoil_dur_min, rise_time=params.dG)

    else:
        spoil_duration = float(spoil_duration)
        spoil_duration = (np.ceil(spoil_duration / scanner_parameters.grad_raster_time)
                          * scanner_parameters.grad_raster_time)

        gx_spoil1 = make_trapezoid(channel='y', system=scanner_parameters, area=gx.area + phase_encoding_area,
                                   duration=spoil_duration, rise_time=params.dG)
        gx_spoil2 = make_trapezoid(channel='y', system=scanner_parameters, area=gx.area - phase_encoding_area,
                                   duration=spoil_duration, rise_time=params.dG)

    # readout gradient with spoilers

    gx_sp1_times = np.array(
        [
            0,
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil1.rise_time) / scanner_parameters.grad_raster_time),
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil1.rise_time + gx_spoil1.flat_time) / scanner_parameters.grad_raster_time),
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil1.rise_time + gx_spoil1.flat_time + gx_spoil1.fall_time) / scanner_parameters.grad_raster_time)
        ]
    )
    gx_sp1_amp = np.array(
        [
            0,
            gx_spoil1.amplitude,
            gx_spoil1.amplitude,
            gx.amplitude
        ]
    )
    gx_sp1 = make_extended_trapezoid(channel='y', system=scanner_parameters, times=gx_sp1_times, amplitudes=gx_sp1_amp)

    gx_sp2_times = np.array(
        [
            0,
            scanner_parameters.grad_raster_time * np.round(
                (gx.flat_time) / scanner_parameters.grad_raster_time)
        ]
    )
    gx_sp2_amp = np.array(
        [
            gx.amplitude,
            gx.amplitude
        ]
    )
    gx_sp2 = make_extended_trapezoid(channel='y', system=scanner_parameters, times=gx_sp2_times, amplitudes=gx_sp2_amp)

    gx_sp3_times = np.array(
        [
            0,
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil2.rise_time) / scanner_parameters.grad_raster_time),
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil2.rise_time + gx_spoil2.flat_time) / scanner_parameters.grad_raster_time),
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil2.rise_time + gx_spoil2.flat_time + gx_spoil2.fall_time) / scanner_parameters.grad_raster_time)
        ]
    )

    gx_sp3_amp = np.array(
        [
            gx.amplitude,
            gx_spoil2.amplitude,
            gx_spoil2.amplitude,
            0
        ]
    )
    gx_sp3 = make_extended_trapezoid(channel='y', system=scanner_parameters, times=gx_sp3_times, amplitudes=gx_sp3_amp)

    if rewind_duration == 'min':
        gx_pre_dur_min = params.dG + (gx.area * 1.5) / scanner_parameters.max_grad
        gx_pre_dur_min = np.ceil(gx_pre_dur_min / scanner_parameters.grad_raster_time) * scanner_parameters.grad_raster_time
        if gx_pre_dur_min < 2 * params.dG:
            gx_pre_dur_min = 2 * params.dG + scanner_parameters.grad_raster_time

        gx_pre = make_trapezoid(channel="y", system=scanner_parameters, area=gx.area * 1.50,
                                duration=gx_pre_dur_min, rise_time=params.dG)

    else:
        rewind_duration = float(rewind_duration)
        rewind_duration = (np.ceil(rewind_duration / scanner_parameters.grad_raster_time)
                          * scanner_parameters.grad_raster_time)
        gx_pre = make_trapezoid(channel="y", system=scanner_parameters, area=gx.area * 1.50,
                                rise_time=params.dG, duration=rewind_duration)

    if united:
        return gx_sp1, gx_sp2, gx_sp3, t_gx, gx_pre
    else:
        return gx_spoil1, gx, gx_spoil2, t_gx, gx_pre


def readout_grad_x(params, scanner_parameters, spoil_duration, rewind_duration,
                   united: bool, phi_radial, phase_encoding_area, **kwargs):
    """
    merged with spoiler gradients readout block in y direction
    phi_radial - angle for radial acquisition
                 (0 - for cartesian)
    phase_encoding_area - additional gradient of phase encoding for spoiler
    @author: petrm
    """
    if 'dixon_delta_te' in kwargs:
        dixon_delta_te = kwargs['dixon_delta_te'] #TODO: add exceptions??
    else:
        dixon_delta_te = 0

    phase_encoding_area = phase_encoding_area * np.sin(phi_radial) * (-1)

    # Create readout gradient
    readout_time = round(1 / params.BW_pixel, 8)
    k_read = np.double(params.Nf) / np.double(params.FoV_f)
    t_gx = np.ceil(readout_time / scanner_parameters.grad_raster_time) * scanner_parameters.grad_raster_time
    gx = make_trapezoid(channel='x', system=scanner_parameters, flat_area=k_read * np.cos(phi_radial),
                        flat_time=t_gx + 2 * scanner_parameters.adc_dead_time, rise_time=params.dG)
    if vars(params).get('Dixon') is not None:
        if params.Dixon == True:
            Dixon_prephasor_area = gx.amplitude * dixon_delta_te
        else:
            Dixon_prephasor_area = 0

    # generate gx spoiler gradient - G_crr
    if spoil_duration == 'min':
        gx_spoil_dur_min = 2 * params.dG + (gx.area + abs(phase_encoding_area) + abs(Dixon_prephasor_area))/scanner_parameters.max_grad
        gx_spoil_dur_min = np.ceil(gx_spoil_dur_min / scanner_parameters.grad_raster_time) * scanner_parameters.grad_raster_time
        if gx_spoil_dur_min < 2 * params.dG:
            gx_spoil_dur_min = 2 * params.dG + scanner_parameters.grad_raster_time



        gx_spoil2 = make_trapezoid(channel='x', system=scanner_parameters, area=gx.area + phase_encoding_area + Dixon_prephasor_area,
                                   rise_time=params.dG, duration=gx_spoil_dur_min)
        gx_spoil1 = make_trapezoid(channel='x', system=scanner_parameters, area=gx.area - phase_encoding_area - Dixon_prephasor_area,
                                   rise_time=params.dG, duration=gx_spoil_dur_min)

    else:
        spoil_duration = float(spoil_duration)
        spoil_duration = (np.ceil(spoil_duration / scanner_parameters.grad_raster_time)
                          * scanner_parameters.grad_raster_time)
        gx_spoil1 = make_trapezoid(channel='x', system=scanner_parameters, area=gx.area + phase_encoding_area + Dixon_prephasor_area,
                                   duration=spoil_duration, rise_time=params.dG)
        gx_spoil2 = make_trapezoid(channel='x', system=scanner_parameters, area=gx.area - phase_encoding_area - Dixon_prephasor_area,
                                   duration=spoil_duration, rise_time=params.dG)

    # readout gradient with spoilers

    gx_sp1_times = np.array(
        [
            0,
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil1.rise_time) / scanner_parameters.grad_raster_time),
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil1.rise_time + gx_spoil1.flat_time) / scanner_parameters.grad_raster_time),
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil1.rise_time + gx_spoil1.flat_time + gx_spoil1.fall_time) / scanner_parameters.grad_raster_time)
        ]
    )
    gx_sp1_amp = np.array(
        [
            0,
            gx_spoil1.amplitude,
            gx_spoil1.amplitude,
            gx.amplitude
        ]
    )
    gx_sp1 = make_extended_trapezoid(channel='x', system=scanner_parameters, times=gx_sp1_times, amplitudes=gx_sp1_amp)

    gx_sp2_times = np.array(
        [
            0,
            scanner_parameters.grad_raster_time * np.round(
                (gx.flat_time) / scanner_parameters.grad_raster_time)
        ]
    )
    gx_sp2_amp = np.array(
        [
            gx.amplitude,
            gx.amplitude
        ]
    )
    gx_sp2 = make_extended_trapezoid(channel='x', system=scanner_parameters, times=gx_sp2_times, amplitudes=gx_sp2_amp)

    gx_sp3_times = np.array(
        [
            0,
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil2.rise_time) / scanner_parameters.grad_raster_time),
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil2.rise_time + gx_spoil2.flat_time) / scanner_parameters.grad_raster_time),
            scanner_parameters.grad_raster_time * np.round(
                (gx_spoil2.rise_time + gx_spoil2.flat_time + gx_spoil2.fall_time) / scanner_parameters.grad_raster_time)
        ]
    )

    gx_sp3_amp = np.array(
        [
            gx.amplitude,
            gx_spoil2.amplitude,
            gx_spoil2.amplitude,
            0
        ]
    )
    gx_sp3 = make_extended_trapezoid(channel='x', system=scanner_parameters, times=gx_sp3_times, amplitudes=gx_sp3_amp)

    if rewind_duration == 'min':

        gx_pre_dur_min = params.dG + (gx.area * 1.5)/scanner_parameters.max_grad
        gx_pre_dur_min = np.ceil(gx_pre_dur_min / scanner_parameters.grad_raster_time) * scanner_parameters.grad_raster_time
        if gx_pre_dur_min < 2 * params.dG:
            gx_pre_dur_min = 2 * params.dG + scanner_parameters.grad_raster_time

        gx_pre = make_trapezoid(channel="x", system=scanner_parameters, area=gx.area * 1.50,
                                duration = gx_pre_dur_min, rise_time=params.dG)

    else:
        rewind_duration = float(rewind_duration)
        rewind_duration = (np.ceil(rewind_duration / scanner_parameters.grad_raster_time)
                          * scanner_parameters.grad_raster_time)
        gx_pre = make_trapezoid(channel="x", system=scanner_parameters, area=gx.area * 1.50,
                                rise_time=params.dG, duration=rewind_duration)

    if united:
        return gx_sp1, gx_sp2, gx_sp3, t_gx, gx_pre
    else:
        return gx_spoil1, gx, gx_spoil2, t_gx, gx_pre