# -*- 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 MRI_sequences.pypulseq.make_sinc_pulse import make_sinc_pulse
from MRI_sequences.pypulseq.make_trapezoid import make_trapezoid
from MRI_sequences.pypulseq.make_extended_trapezoid import make_extended_trapezoid
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=0.5,
        time_bw_product=round(params['t_BW_product_ref'], 8),
        phase_offset=rf180_phase,
        use="refocusing",
        return_gz=True,
    )

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

    # spoil gradient around 180 RF pulse - G_crs
    # t_gz_spoil = (np.ceil(params['t_ref'] / 2 / scanner_parameters.grad_raster_time)
    #               * scanner_parameters.grad_raster_time)

    if spoil_duration == 'min':
        gz_spoil1 = make_trapezoid(channel='z', system=scanner_parameters, area= gz90_area,
                                   rise_time=params['dG'], flat_time=params['dG'])
        gz_spoil2 = make_trapezoid(channel='z', system=scanner_parameters, area= gz90_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
        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(params, scanner_parameters, spoil_duration, united: bool):
    # 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,
                        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,
                                  flat_time=params['dG'], 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_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 united:
        return gx_sp1, gx_sp2, gx_sp3, t_gx, gx.area
    else:
        return gx_spoil, gx, gx_spoil, t_gx, gx.area