# -*- coding: utf-8 -*-
"""
Created on Fri Apr  5 10:43:09 2024

@author: zilya
"""

import numpy as np
from LF_scanner.pypulseq.Sequence.sequence import Sequence
from LF_scanner.pypulseq.calc_duration import calc_duration
from LF_scanner.pypulseq.check_timing import check_timing
from LF_scanner.pypulseq.make_adc import make_adc
from LF_scanner.pypulseq.make_delay import make_delay
from LF_scanner.pypulseq.make_sinc_pulse import make_sinc_pulse
from LF_scanner.pypulseq.make_trapezoid import make_trapezoid
from LF_scanner.pypulseq.opts import Opts


def seqgen_STEAM(param):
    
    
    # Read scanner parameters from the params structure
    scanner_parameters = Opts(max_grad = param.G_amp_max, grad_unit = 'Hz/m',\
                              max_slew = param.G_slew_max, slew_unit = 'Hz/m/s',\
                              grad_raster_time = param.grad_raster_time,\
                              rf_raster_time = param.rf_raster_time,\
                              adc_raster_time = 1/(param.BW_per_point*param.N_point),\
                              block_duration_raster = max(param.grad_raster_time,param.rf_raster_time))
    
    # For all modules
    gradient_ramp_time = param.grad_raster_time*np.ceil((scanner_parameters.max_grad/scanner_parameters.max_slew)/param.grad_raster_time)
        
    # Calculate slice-selection gradient
    pulse_bandwidth = np.double(param.t_BW_product_ex) / np.double(param.t_ex) #In Hz
    gradient_voxel_amp = np.double(pulse_bandwidth) / np.double(param.voxel_thkn) #In Hz/m
    gradient_voxel = make_trapezoid(channel = 'z',
                                    flat_time = np.double(param.t_ex), 
                                    rise_time = gradient_ramp_time, 
                                    amplitude = gradient_voxel_amp, 
                                    system = scanner_parameters)
    # Generate the RF pulse
    exc_pulse = make_sinc_pulse(flip_angle = np.radians(param.FA),\
                                duration = np.double(param.t_ex),\
                                phase_offset = 0,\
                                time_bw_product = param.t_BW_product_ex,\
                                delay = gradient_ramp_time,\
                                system = scanner_parameters)
   
    # Generate the ADC event
    adc_module = make_adc(num_samples = param.N_point, duration = 0.5/param.BW_per_point, delay = gradient_ramp_time, system = scanner_parameters)
    
    # Calculate the TE pause time
    tau_delay1 = param.tau1 - calc_duration(gradient_voxel)
    tau_delay1 = param.grad_raster_time*np.floor(tau_delay1/param.grad_raster_time)
    
    tau_delay2 = param.tau2 - calc_duration(gradient_voxel)
    tau_delay2 = param.grad_raster_time*np.floor(tau_delay2/param.grad_raster_time)
    
    tau_delay3 = param.tau1 - 0.5*calc_duration(gradient_voxel)
    tau_delay3 = param.grad_raster_time*np.floor(tau_delay3/param.grad_raster_time)
    
    tau_delay1 = make_delay(tau_delay1) if tau_delay1 > 0 else None
    tau_delay2 = make_delay(tau_delay2) if tau_delay2 > 0 else None
    tau_delay3 = make_delay(tau_delay3) if tau_delay3 > 0 else None
    

    # Calculate the TR fillers
    tr_delay = param.TR - (2*param.tau1 + param.tau2 + calc_duration(adc_module) + 0.5*calc_duration(gradient_voxel))
    tr_delay = param.grad_raster_time*np.floor(tr_delay/param.grad_raster_time)
    tr_delay = make_delay(tr_delay) if tr_delay > 0 else None
    
    seq = Sequence(system = scanner_parameters)
    phase_shift = 0
    # Dummy scans
    for n in range(int(param.average)):
        gradient_voxel.channel = 'z'
        seq.add_block(exc_pulse,gradient_voxel)
        if tau_delay1 is not None:
            seq.add_block(tau_delay1)
        gradient_voxel.channel = 'y'
        seq.add_block(exc_pulse,gradient_voxel)
        if tau_delay2 is not None:
            seq.add_block(tau_delay2)
        gradient_voxel.channel = 'x'
        seq.add_block(exc_pulse,gradient_voxel)
        if tau_delay3 is not None:
            seq.add_block(tau_delay3)
        seq.add_block(adc_module)
        if tr_delay is not None:
            seq.add_block(tr_delay)
    print(check_timing(seq)) 
    return seq