# -*- coding: utf-8 -*-
"""
A subroutine to add IR block.
Requires the params structure as input.

@author: petrm
"""

import os
import sys

# need to add the project dir with the utilities folder to path: 
project_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
sys.path.append(project_dir)

from pypulseq.make_sinc_pulse import make_sinc_pulse
from pypulseq.make_trapezoid import make_trapezoid
from pypulseq.make_delay import make_delay
from pypulseq.calc_duration import calc_duration as calc_dur

from math import pi
import numpy as np

def IR_block(params, hw, IR_geom_offset):

    # function creates inversion recovery block with delay
    
# =============================================================================
# Get params from GUI
# =============================================================================

    t_inv = params['t_inv']
    t_BW_product_inv = params['t_BW_product_inv']
    TI_calc_block = params['TI_calc_block']

    if not 'dG' in params:
        params['dG'] = hw.grad_ramp_time

    spoiler_area = params['magn_prep_spoiler_area']
        
    if not 'flip_IR' in params:
        params['flip_IR'] = round(180 * pi / 180, 3)


# =============================================================================
# We agreed that IR in 3D is always nonselective. FoV_ss is a 3D indicator
# =============================================================================
    if 'FoV_ss' in params:
        slice_selective = False
    else:
        slice_selective = True

# =============================================================================
# Calculate module
# =============================================================================

    rf_IR_output = make_sinc_pulse(flip_angle = params['flip_IR'], 
                                      system = hw, 
                                      duration = t_inv,
                                      slice_thickness = params['inv_thkn'],
                                      apodization = 0.5,
                                      time_bw_product = round(t_BW_product_inv, 8), 
                                      freq_offset = 0,
                                      phase_offset = 90 * pi / 180,
                                      return_gz = slice_selective)
    
    
    if slice_selective:   # in this case, rf_IR_output is a list of namespaces
        rf_IR = rf_IR_output[0]
        gz_IR = rf_IR_output[1] 
        rf_IR.freq_offset = IR_geom_offset * gz_IR.amplitude # calculate IR freq offset from the geom offset (for slices)
        duration_IR_object = calc_dur(gz_IR)
        
        IR_return_objects_list = [rf_IR, gz_IR]
        
    else:                 # in this case, rf_IR_output is actually rf_IR itself
        rf_IR = rf_IR_output 
        duration_IR_object = calc_dur(rf_IR)
        IR_return_objects_list = [rf_IR]
        
    # calculate spoiler
        
    gz_IR_spoil = make_trapezoid(channel = "z", 
                                 system = hw, 
                                 area = spoiler_area,
                                 rise_time = params['dG'])
    
    duration_gz_IR_spoil = calc_dur(gz_IR_spoil)

    # calculate delay
    
    TI_delay = TI_calc_block - 1/2*duration_IR_object - duration_gz_IR_spoil
    TI_delay = np.ceil(TI_delay / hw.grad_raster_time) * hw.grad_raster_time 

    delay_IR = make_delay(TI_delay)

    # calculate duration
    
    duration_IR_total = duration_IR_object + duration_gz_IR_spoil + TI_delay
    
    return (IR_return_objects_list, [gz_IR_spoil], [delay_IR]), duration_IR_total