# -*- coding: utf-8 -*-
"""
sequence_generation
______________________

Version: 1.0
Authors: Davis Garrad (Paul Scherrer Institute, CH)
______________________

Wrapper for the API I wrote to generate pulse sequences programmatically in TNMR (Tecmag), with a little more polish than the original driver code.
"""
import frappy_psi.tnmr.sequence_fileformat as se

from pydantic.utils import deep_update

import json

def get_single_pulse_block(name, pulse_width, pulse_height, delay_time, phase_cycle='0'):
    '''Generates a single block of data to create a sequence with.

    Parameters
    ----------
        name: str, just the prefix for the column names. Ensure this is unique over the whole sequence!
        pulse_width: str, in the format '10u' (10 microseconds)
        pulse_height: str, in the format '40' (I'm not honestly sure what units this is in...)
        delay_time: str, in the format '10u'. Minimum value of '0.1u'. If '0', then only the pulse is added (not the delay afterwards). If no unit is given, assumes microseconds.
        phase_cycle: a given phase cycle (steps of 4, 0-3 inc.) (eg., '0 0 1 1 2 3 0 1', etc.) (default '0')

    Returns
    -------
        a dictionary which can be updated with others (with combine_blocks) to generate a larger, more complex sequence.
    '''
    
    if(delay_time.strip()[-1] == 'u'):
        delay_time = float(delay_time.strip()[:-1])
    elif(delay_time.strip()[-1] == 'n'):
        delay_time = float(delay_time.strip()[:-1]) * 1000
    elif(delay_time.strip()[-1] == 'm'):
        delay_time = float(delay_time.strip()[:-1]) / 1e3
    elif(delay_time.strip()[-1] == 's'):
        delay_time = float(delay_time.strip()[:-1]) / 1e6
    else:
        delay_time = float(delay_time.strip()) # assume in us
        
    delay_time_rounded = int(delay_time*10) / 10 # nearest 10ns
    delay_time = delay_time_rounded

    ph = name + '_phase'
    rl = name + '_delay'
    block = se.generate_default_sequence([ ph, rl ] if delay_time > 0 else [ ph ], [ pulse_width, str(delay_time) + 'u' ] if delay_time > 0 else [pulse_width])

    # COLUMNNS
    # PH column
    block['columns'][ph]['F1_Ampl']['value'] = str(pulse_height)
    block['columns'][ph]['Delay'] = str(pulse_width)
    block['columns'][ph]['F1_UnBlank']['value'] = '1'
    block['columns'][ph]['Rx_Blank']['value'] = '1'

    if(delay_time > 0):
        # delay column
        block['columns'][rl]['F1_UnBlank']['value'] = '1'
        block['columns'][rl]['Rx_Blank']['value'] = '1'

    if(phase_cycle != ''):
        table_name = f'ph_{name}'
        block['columns'][ph]['F1_Ph']['table'] = table_name
        block['tables'][table_name] = { 'values': phase_cycle, 'typestr': 'HP', 'start': 1 }

    return block

def get_initial_block():
    '''Generates the mandatory initial block, which consists of phase reset and unblanking
    
    Returns
    -------
        a dictionary which can be updated with others (with combine_blocks) to generate a larger, more complex sequence.
    '''
    block = se.generate_default_sequence(['Phase reset', 'Unblank'], ['1u', '10u'])
    # Phase reset
    block['columns']['Phase reset']['F1_PhRst']['value'] = '1'
    # Unblank
    block['columns']['Unblank']['F1_UnBlank']['value'] = '1'
    block['columns']['Unblank']['Rx_Blank']['value'] = '1'

    return block

def get_final_block(ringdown_time, preacquire_time, acquire_time, cooldown_time, acq_phase_cycle='0', num_acq_points=1024):
    '''Generates the final block of data to create a sequence with.

    Parameters
    ----------
        ringdown_time: str, of format '10u', how long to ringdown
        preacquire_time:  str, in the format '10u', how long to wait after ringdown before acquiring data. A good default is 1 usec
        acquire_time: str, in the format '10u', how long to acquire data for
        cooldown_time: str, in the format '40u', how long to wait after acquisition.
        acq_phase_cycle: str, the phase cycle that the acquisition should follow

    Returns
    -------
        a dictionary which can be updated with others to generate a larger, more complex sequence.
    '''
    block = se.generate_default_sequence(['Ringdown', 'RX On', 'Acquisition', 'Finish', ''], [ringdown_time, preacquire_time, acquire_time, cooldown_time, '1u'])
    block['num_acq_points'] = num_acq_points
    
    # ringdown
    block['columns']['Ringdown']['Rx_Blank']['value'] = '1'

    # Acquire
    block['columns']['Acquisition']['Acq']['value'] = '1'
    #block['columns']['Acquisition']['Acq_phase']['value'] = '0'

    if(acq_phase_cycle != ''):
        block['columns']['Acquisition']['Acq_phase']['table'] = 'phacq'
        block['tables']['phacq'] = { 'values': acq_phase_cycle, 'typestr': 'HP', 'start': 1 }
    
    return block

def combine_blocks(l, r):
    '''Combines two dictionaries (just deep copies them, as they each have recursive dictionaries'''
    return deep_update(l, r)

def save_sequence(filename, sequence):
    '''Saves the given sequence to a file'''
    se.create_sequence_file(filename, sequence)

def save_sequence_cfg(filename, sequence):
    '''Saves the sequence to a file in JSON form to be read easier. Appends the extension ".cfg" '''
    with open(filename + '.cfg', 'w') as file:
        json.dump(sequence, file, indent=4)
        file.close()