''' Class to convert Mimosa26 raw data.

General structure of Mimosa26 raw data (32 bit words):
 - First 8 bits are always 0x20 (Mimosa26 HEADER)
 - Next 4 bits are plane number 1 - 6 (plane identifier)
 - Next 2 bits always zero
 - Next two bits contain: data loss flag and frame start flag
 - Rest of 16 bits are actual data words

The raw data structure of Mimosa26 data looks as follows:
 - Frame header HIGH and LOW (contains timestamp, generated from R/O) [word index 0 + 1]
 - Frame number HIGH and LOW (frame number of Mimosa26) [word index 2 + 3]
 - Frame length HIGH and LOW (number of Mimosa26) [word index 4 + 5]
 - Hit data (column and row of hit pixel)
 - ...
 - ...
 - Frame trailer HIGH and LOW (indicates end of Mimosa26 frame) [word index 6 + 7]

'''
import numba
from numba import njit
import numpy as np

MIMOSA_FRAME_CYCLE = 115.2  # us
MIMOSA_FREQ = 40  # MHz
N_ROWS_MIMOSA = 576  # Number of rows
FRAME_UNIT_CYCLE = int(MIMOSA_FRAME_CYCLE * MIMOSA_FREQ)  # = 4608, time for one frame in units of 40 MHz clock cylces
ROW_UNIT_CYCLE = int(MIMOSA_FRAME_CYCLE * MIMOSA_FREQ / N_ROWS_MIMOSA)  # = 8, time to read one row in units of 40 MHz clock cycles
TIMING_OFFSET = -112  # Correct for offset between M26 40 MHz clock and 40 MHz from R/O system. Offset determined by maximum correlation between the time reference and Mimosa26 telescope.
MAX_BUFFER_TIME_SLIP = 5  # max. time (in seconds) for storing hits in buffer before they get removed if no trigger appears
DEFAULT_PYMOSA_M26_HEADER_IDS = [1, 2, 3, 4, 5, 6]  # Default header IDs for the Mimosa26 data generated by the pymosa software. The header IDs are set in the pymosa readout software.

hits_dtype = np.dtype([
    ('plane', '<u1'),
    ('event_number', '<i8'),
    ('trigger_number', '<i8'),
    ('trigger_time_stamp', '<i8'),
    ('row_time_stamp', '<i8'),
    ('frame_id', '<i8'),
    ('column', '<u2'),
    ('row', '<u2'),
    ('event_status', '<u4')])

telescope_data_dtype = np.dtype([
    ('plane', '<u1'),
    ('time_stamp', '<i8'),
    ('frame_id', '<i8'),
    ('column', '<u2'),
    ('row', '<u2'),
    ('frame_status', '<u4')])

trigger_data_dtype = np.dtype([
    ('event_number', '<i8'),
    ('trigger_number', '<i8'),
    ('trigger_time_stamp', '<i8'),
    ('trigger_status', '<u4')])

# Error codes
TRIGGER_NUMBER_ERROR = 0x00000001  # Trigger number has not increased by one
NO_TRIGGER_WORD_ERROR = 0x00000002  # Event has no trigger word associated
TRIGGER_TIMESTAMP_OVERFLOW = 0x00000004  # Indicating the overflow of the trigger timestamp
TRIGGER_NUMBER_OVERFLOW = 0x00000008  # Indicating the overflow of the trigger number
DATA_ERROR = 0x00000010  # Indicating any occurrence of data errors in the Momosa26 protocol (e.g., invalid column/row, invalid data length, data loss)
TIMESTAMP_OVERFLOW = 0x00000020  # Indicating the overflow of the Mimosa26 timestamp
FRAME_ID_OVERFLOW = 0x00000040  # Indicating the overflow of the Mimosa26 frame ID
OVERFLOW_FLAG = 0x00000080  # Indicating the occurrence of the overflow flag for a particular Mimosa26 row

# Mimosa26 raw data
@njit
def is_mimosa_data(word):  # Check for Mimosa data word
    return (0xff000000 & word) == 0x20000000

@njit
def get_plane_number(word):  # There are 6 planes in the stream, starting from 1; return plane number
    return (word >> 20) & 0xf

# Frame header
@njit
def is_frame_header(word):  # Check if frame header high word (frame start flag is set by R/0)
    return (0x00010000 & word) == 0x00010000

@njit
def is_data_loss(word):  # Indicates data loss
    return (0x00020000 & word) == 0x00020000

@njit
def get_m26_timestamp_low(word):  # Timestamp of Mimosa26 data from frame header low (generated by R/0)
    return 0x0000ffff & word

@njit
def get_m26_timestamp_high(word):  # Timestamp of Mimosa26 data from frame header high (generated by R/0)
    return (0x0000ffff & word) << 16

@njit
def is_frame_header0(word):  # Check if frame header0 word
    return (0x0000ffff & word) == 0x00005555

@njit
def is_frame_header1(word, plane):  # Check if frame header1 word for the actual plane
    return (0x0000ffff & word) == (0x00005550 | plane)

# Frame counter
@njit
def get_frame_id_low(word):  # Get the frame id from the frame id low word
    return 0x0000ffff & word

@njit
def get_frame_id_high(word):  # Get the frame id from the frame id high word
    return (0x0000ffff & word) << 16

# Data length
@njit
def get_frame_length(word):  # Get length of Mimosa26 frame
    return (0x0000ffff & word)

# Status / line word
@njit
def get_n_words(word):  # Return the number of data words for the actual row
    return 0x0000000f & word

@njit
def get_row(word):  # Extract row from Mimosa26 hit word
    return (0x00007ff0 & word) >> 4

@njit
def has_overflow(word):
    return (0x00008000 & word) != 0

# State word
@njit
def get_n_hits(word):  # Returns the number of hits given by actual column word
    return 0x00000003 & word

@njit
def get_column(word):  # Extract column from Mimosa26 hit word
    return (0x00001ffc & word) >> 2

# Frame trailer
@njit
def is_frame_trailer0(word):  # Check if frame trailer0 word
    return (0x0000ffff & word) == 0xaa50

@njit
def is_frame_trailer1(word, plane):  # Check if frame trailer1 word for the actual plane
    return (0x0000ffff & word) == (0xaa50 | plane)

# Trigger words
@njit
def is_trigger_word(word):  # Check if TLU word (trigger)
    return (0x80000000 & word) == 0x80000000

@njit
def get_trigger_timestamp(word):  # Get timestamp of TLU word
    return (word & 0x7fff0000) >> 16

@njit
def get_trigger_number(word, trigger_data_format):  # Get trigger number of TLU word
    if trigger_data_format == 2:
        return word & 0x0000ffff
    else:
        return word & 0x7fffffff

class RawDataInterpreter(object):
    ''' Class to convert the raw data chunks to hits'''

    def __init__(self, analyze_m26_header_ids=None):
        '''
        Parameters:
        -----------
        analyze_m26_header_ids : list
            List of Mimosa26 header IDs that will be interpreted.
            If None, the value defaults to the global value raw_data_interpreter.DEFAULT_PYMOSA_M26_HEADER_IDS.
        '''
        if analyze_m26_header_ids is None:
            self.analyze_m26_header_ids = DEFAULT_PYMOSA_M26_HEADER_IDS
        else:
            self.analyze_m26_header_ids = analyze_m26_header_ids
        for analyze_m26_header_id in self.analyze_m26_header_ids:
            if analyze_m26_header_id < 0 or analyze_m26_header_id >= 2**16:
                raise ValueError('Invalid header ID.')
        self.analyze_m26_header_ids = np.asarray(self.analyze_m26_header_ids, dtype=np.uint16)
        self.plane_id_to_index = -1 * np.ones(shape=max(self.analyze_m26_header_ids) + 1, dtype=np.int32)
        for plane_index, plane_id in enumerate(self.analyze_m26_header_ids):
            self.plane_id_to_index[plane_id] = plane_index
        self.reset()

    def reset(self):  # Reset variables
        # Temporary arrays
        self.trigger_data = np.zeros(shape=0, dtype=trigger_data_dtype)
        self.trigger_data_index = np.int64(-1)
        self.telescope_data = np.zeros(shape=0, dtype=telescope_data_dtype)
        self.telescope_data_index = np.int64(-1)

        # Raw data interpreter
        # Per frame variables
        self.m26_frame_ids = np.zeros(shape=len(self.analyze_m26_header_ids), dtype=np.int64)  # The Mimosa26 frame ID of the actual frame
        self.m26_frame_length = np.zeros(shape=len(self.analyze_m26_header_ids), dtype=np.uint32)  # The number of "useful" data words for the actual frame
        self.m26_data_loss = np.ones(len(self.analyze_m26_header_ids), dtype=np.bool_)  # The data loss status for the actual frame
        self.m26_word_index = np.zeros(shape=len(self.analyze_m26_header_ids), dtype=np.uint32)  # The word index per device of the actual frame
        self.m26_timestamps = np.zeros(shape=len(self.analyze_m26_header_ids), dtype=np.int64)  # The timestamp for each plane (in units of 40 MHz)
        self.last_m26_timestamps = np.zeros(shape=len(self.analyze_m26_header_ids), dtype=np.int64)
        self.m26_n_words = np.zeros(shape=len(self.analyze_m26_header_ids), dtype=np.uint32)  # The number of words containing column / row info
        self.m26_rows = np.zeros(shape=len(self.analyze_m26_header_ids), dtype=np.uint32)  # The actual readout row (rolling shutter)
        self.m26_frame_status = np.zeros(shape=len(self.analyze_m26_header_ids), dtype=np.uint32)  # The status flags for the actual frames
        self.last_completed_m26_frame_ids = -1 * np.ones(shape=len(self.analyze_m26_header_ids), dtype=np.int64)  # The status if the frame is complete for the actual frame
        # Per event variables
        self.event_number = np.int64(-1)  # The event number of the actual trigger, event number starts at 0
        self.trigger_number = np.int64(-1)  # The trigger number of the actual trigger
        self.trigger_timestamp = np.int64(0)  # The trigger timestamp of the actual trigger

        # Event builder
        self.hits = np.zeros(shape=0, dtype=hits_dtype)
        self.hits_index = np.int64(-1)

        # Properties
        self._add_missing_events = False
        self._timing_offset = TIMING_OFFSET

    @property
    def add_missing_events(self):
        return self._add_missing_events

    @add_missing_events.setter
    def add_missing_events(self, value):
        self._add_missing_events = bool(value)

    @property
    def timing_offset(self):
        return self._timing_offset

    @timing_offset.setter
    def timing_offset(self, value):
        self._timing_offset = int(value)

    def interpret_raw_data(self, raw_data=None, build_all_events=False):
        ''' Converting the raw data array to a hit array.
        The is the only function that needs to be called to convert the raw data.

        Parameters:
        -----------
        raw_data : np.array
            The array with the raw data words.
        build_all_events : bool
            If True, build all events from the remaining trigger_data and telescope_data_array.
            Use this only after the last raw data chunk to receive the the remaining events in the buffers.
        '''
        if raw_data is None:
            raw_data = np.zeros(shape=0, dtype=np.uint32)
        if self.telescope_data_index != -1:
            telescope_data_index_start = self.telescope_data_index + 1
        else:
            telescope_data_index_start = 0
        # Analyze raw data
        self.trigger_data, self.trigger_data_index, self.telescope_data, self.telescope_data_index, self.m26_frame_ids, self.m26_frame_length, self.m26_data_loss, self.m26_word_index, self.m26_timestamps, self.last_m26_timestamps, self.m26_n_words, self.m26_rows, self.m26_frame_status, self.last_completed_m26_frame_ids, self.event_number, self.trigger_number, self.trigger_timestamp = _interpret_raw_data(
            raw_data=raw_data,
            trigger_data=self.trigger_data,
            trigger_data_index=self.trigger_data_index,
            telescope_data=self.telescope_data,
            telescope_data_index=self.telescope_data_index,
            m26_frame_ids=self.m26_frame_ids,
            m26_frame_length=self.m26_frame_length,
            m26_data_loss=self.m26_data_loss,
            m26_word_index=self.m26_word_index,
            m26_timestamps=self.m26_timestamps,
            last_m26_timestamps=self.last_m26_timestamps,
            m26_n_words=self.m26_n_words,
            m26_rows=self.m26_rows,
            m26_frame_status=self.m26_frame_status,
            last_completed_m26_frame_ids=self.last_completed_m26_frame_ids,
            event_number=self.event_number,
            trigger_number=self.trigger_number,
            trigger_timestamp=self.trigger_timestamp,
            add_missing_events=self.add_missing_events,
            build_all_events=build_all_events,
            analyze_m26_header_ids=self.analyze_m26_header_ids,
            plane_id_to_index=self.plane_id_to_index)

        # Get data from telescope (just hit data, no assignment to events or data multiplication)
        telescope_data = self.telescope_data[telescope_data_index_start:self.telescope_data_index + 1].copy()

        # Build events
        self.trigger_data, self.trigger_data_index, self.telescope_data, self.telescope_data_index, self.hits, self.hits_index = _build_events(
            trigger_data=self.trigger_data,
            trigger_data_index=self.trigger_data_index,
            telescope_data=self.telescope_data,
            telescope_data_index=self.telescope_data_index,
            hits=self.hits,
            hits_index=self.hits_index,
            last_completed_m26_frame_ids=self.last_completed_m26_frame_ids,
            timing_offset=self.timing_offset,
            build_all_events=build_all_events,
            analyze_m26_header_ids=self.analyze_m26_header_ids,
            plane_id_to_index=self.plane_id_to_index)
        # Create a copy of the hits array that is returned
        hits = self.hits[:self.hits_index + 1].copy()
        self.hits_index -= (self.hits_index + 1)

        return hits, telescope_data

@njit(locals={'trigger_data_index': numba.int64, 'telescope_data_index': numba.int64, 'trigger_status': numba.uint32, 'last_trigger_number': numba.int64, 'last_trigger_timestamp': numba.int64, 'n_missing_events': numba.uint32})
def _interpret_raw_data(raw_data, trigger_data, trigger_data_index, telescope_data, telescope_data_index, m26_frame_ids, m26_frame_length, m26_data_loss, m26_word_index, m26_timestamps, last_m26_timestamps, m26_n_words, m26_rows, m26_frame_status, last_completed_m26_frame_ids, event_number, trigger_number, trigger_timestamp, add_missing_events, build_all_events, analyze_m26_header_ids, plane_id_to_index):
    ''' This function is interpreting the Mimosa26 telescope raw data and creates temporary trigger and telescope data arrays.
    The interpreter checks for trigger and Mimosa26 data errors.

    Parameters:
    -----------
    raw_data : np.array
        The array with the raw data words.
    TBD
    '''
    # Loop over the raw data words
    for raw_data_word in raw_data:
        if is_mimosa_data(raw_data_word):  # Check if word is from Mimosa26.
            # Check to which plane the data belongs
            plane_id = get_plane_number(raw_data_word)  # The actual_plane if the actual word belongs to (0 to 5)
            for analyze_m26_header_id in analyze_m26_header_ids:
                if plane_id == analyze_m26_header_id:
                    break
            else:
                continue  # Do not interpret data of planes which should be skipped
            plane_index = plane_id_to_index[plane_id]
            # In the following, interpretation of the raw data words of the actual plane
            # Check for data loss bit set by the M26 RX FSM
            if is_data_loss(raw_data_word):
                # Setting the data loss flag to true.
                # The data loss bit is set by the M26 RX FSM.
                # The bit is set only once after each data loss, i.e.,
                # the first data word after the lost data words.
                m26_data_loss[plane_index] = True
            if is_frame_header(raw_data_word):  # New frame for actual plane, M26 timestamp (LSB), frame header0
                # Get Mimosa26 timestamp from raw data word (LSB)
                last_m26_timestamps[plane_index] = m26_timestamps[plane_index]
                m26_timestamps[plane_index] = (m26_timestamps[plane_index] & 0x7fffffffffff0000) | get_m26_timestamp_low(raw_data_word)
                m26_word_index[plane_index] = 0
                # Reset parameters after header
                m26_frame_length[plane_index] = 0
                m26_n_words[plane_index] = 0
                # Set the status bits for priviously incomplete frames
                index = telescope_data_index
                while index >= 0:
                    if telescope_data[index]['plane'] == plane_id:
                        if telescope_data[index]['frame_id'] > last_completed_m26_frame_ids[plane_index]:
                            telescope_data[index]['frame_status'] |= DATA_ERROR
                        else:
                            break
                    index -= 1
                m26_data_loss[plane_index] = False
                m26_frame_status[plane_index] = 0
            elif m26_data_loss[plane_index] is True:  # Trash data
                # Nothing to do, do not trust data
                continue
            else:  # Interpreting M26 raw data
                m26_word_index[plane_index] += 1
                if m26_word_index[plane_index] == 1:  # Mimosa26 timestamp, M26 timestamp (MSB), frame header1
                    # Check for 32bit timestamp overflow
                    if m26_timestamps[plane_index] >= 0 and get_m26_timestamp_high(raw_data_word) < (m26_timestamps[plane_index] & 0x00000000ffff0000):
                        m26_frame_status[plane_index] |= TIMESTAMP_OVERFLOW
                        m26_timestamps[plane_index] = np.int64(2**32) + m26_timestamps[plane_index]
                    # Get Mimosa26 timestamp from raw data word (MSB)
                    m26_timestamps[plane_index] = get_m26_timestamp_high(raw_data_word) | (m26_timestamps[plane_index] & 0x7fffffff0000ffff)
                elif m26_word_index[plane_index] == 2:  # Mimosa26 frame ID
                    # Get Mimosa26 frame ID from raw data word (LSB)
                    m26_frame_ids[plane_index] = (m26_frame_ids[plane_index] & 0x7fffffffffff0000) | get_frame_id_low(raw_data_word)
                elif m26_word_index[plane_index] == 3:  # Mimosa26 frame ID
                    # Check for 32bit frame ID overflow
                    if m26_frame_ids[plane_index] >= 0 and get_frame_id_high(raw_data_word) < (m26_frame_ids[plane_index] & 0x00000000ffff0000):
                        m26_frame_status[plane_index] |= FRAME_ID_OVERFLOW
                        m26_frame_ids[plane_index] = np.int64(2**32) + m26_frame_ids[plane_index]
                    # Get Mimosa26 frame ID from raw data word (MSB)
                    m26_frame_ids[plane_index] = get_frame_id_high(raw_data_word) | (m26_frame_ids[plane_index] & 0x7fffffff0000ffff)
                elif m26_word_index[plane_index] == 4:  # Mimosa26 frame length
                    m26_frame_length[plane_index] = get_frame_length(raw_data_word)
                    if m26_frame_length[plane_index] > 570:  # Defined in the Mimosa26 protocol, no more than 570 "useful" data words
                        m26_data_loss[plane_index] = True
                        continue
                elif m26_word_index[plane_index] == 5:  # Mimosa26 frame length, a second time
                    if m26_frame_length[plane_index] != get_frame_length(raw_data_word):  # DO0 & DO1 should always have the same data length
                        m26_data_loss[plane_index] = True
                        continue
                    else:
                        m26_frame_length[plane_index] += get_frame_length(raw_data_word)
                elif m26_word_index[plane_index] == 5 + m26_frame_length[plane_index] + 1:  # Frame trailer0
                    if not is_frame_trailer0(raw_data_word):
                        m26_data_loss[plane_index] = True
                        continue
                elif m26_word_index[plane_index] == 5 + m26_frame_length[plane_index] + 2:  # Frame trailer1
                    if not is_frame_trailer1(raw_data_word, plane=plane_id):
                        m26_data_loss[plane_index] = True
                        continue
                    else:
                        last_completed_m26_frame_ids[plane_index] = m26_frame_ids[plane_index]
                elif m26_word_index[plane_index] > 5 + m26_frame_length[plane_index] + 2:  # Ignore any occurrence of additional raw data words
                    m26_data_loss[plane_index] = True
                    continue
                else:  # Column / Row words (actual data word with hits)
                    if m26_n_words[plane_index] == 0:  # First word contains the row info and the number of data words for this row
                        if m26_word_index[plane_index] == 5 + m26_frame_length[plane_index]:  # Always even amount of words or this fill word is used
                            # Ignore this fill word
                            continue
                        else:
                            m26_n_words[plane_index] = get_n_words(raw_data_word)
                            m26_rows[plane_index] = get_row(raw_data_word)  # Get row from data word
                            if m26_rows[plane_index] >= 576:  # Row overflow
                                m26_data_loss[plane_index] = True
                                continue
                        if has_overflow(raw_data_word):
                            m26_frame_status[plane_index] |= OVERFLOW_FLAG  # set overflow bit
                        else:
                            m26_frame_status[plane_index] & ~OVERFLOW_FLAG  # unset overflow bit
                    else:
                        m26_n_words[plane_index] = m26_n_words[plane_index] - 1  # Count down the words
                        n_hits = get_n_hits(raw_data_word)
                        column = get_column(raw_data_word)  # Get column from data word
                        if column >= 1152:  # Column overflow
                            m26_data_loss[plane_index] = True
                            continue
                        for k in range(n_hits + 1):
                            if column + k >= 1152:
                                m26_data_loss[plane_index] = True
                                break
                            # Increase index
                            telescope_data_index += 1
                            # extend telescope data array if neccessary
                            if telescope_data_index >= telescope_data.shape[0]:
                                # remove old hit data from array for each plane individually. Prevents the case that telescope data array gets too big in case
                                # time until next trigger is very large, since telescope data has to be buffered until next trigger.
                                select = (telescope_data['plane'] == plane_id)
                                select &= (telescope_data['time_stamp'] < (m26_timestamps[plane_index] - MAX_BUFFER_TIME_SLIP * MIMOSA_FREQ * 10**6))
                                count_outdated = np.sum(select)
                                if count_outdated:
                                    telescope_data = telescope_data[~select]
                                    telescope_data_index = telescope_data_index - count_outdated
                                # extend telescope data array if neccessary
                                telescope_data_tmp = np.zeros(shape=max(1, int(raw_data.shape[0] / 2)), dtype=telescope_data_dtype)
                                telescope_data = np.concatenate((telescope_data, telescope_data_tmp))

                            # Store hits
                            telescope_data[telescope_data_index]['plane'] = plane_id
                            telescope_data[telescope_data_index]['time_stamp'] = m26_timestamps[plane_index]
                            telescope_data[telescope_data_index]['frame_id'] = m26_frame_ids[plane_index]
                            telescope_data[telescope_data_index]['column'] = column + k
                            telescope_data[telescope_data_index]['row'] = m26_rows[plane_index]
                            telescope_data[telescope_data_index]['frame_status'] = m26_frame_status[plane_index]
        elif is_trigger_word(raw_data_word):  # Raw data word is TLU/trigger word
            # Reset trigger status
            trigger_status = 0
            # Get latest telescope timestamp and set trigger timestamp
            last_trigger_timestamp = trigger_timestamp
            # Get largest M26 timestamp
            for tmp_plane_index, _ in enumerate(analyze_m26_header_ids):
                if last_m26_timestamps[tmp_plane_index] > trigger_timestamp:
                    trigger_timestamp = last_m26_timestamps[tmp_plane_index]
            # Calculating 63bit timestamp from 15bit trigger timestamp
            # and last telescope timestamp (frame header timestamp).
            # Assumption: the telescope timestamp is updated more frequent than
            # the 15bit trigger timestamp can overflow. The frame is occurring
            # every 4608 clock cycles (115.2 us).
            # Get trigger timestamp from raw data word
            trigger_timestamp = (0x7fffffffffff8000 & trigger_timestamp) | get_trigger_timestamp(raw_data_word)
            # Check for 15bit trigger timestamp overflow
            if last_trigger_timestamp >= 0 and trigger_timestamp <= last_trigger_timestamp:
                trigger_status |= TRIGGER_TIMESTAMP_OVERFLOW
                trigger_timestamp = np.int64(2**15) + trigger_timestamp
            # Copy of trigger number
            last_trigger_number = trigger_number
            # Check for 16bit trigger number overflow
            if trigger_number >= 0 and get_trigger_number(raw_data_word, trigger_data_format=2) <= (trigger_number & 0x000000000000ffff):
                trigger_status |= TRIGGER_NUMBER_OVERFLOW
                trigger_number = np.int64(2**16) + trigger_number
            # Get trigger number from raw data word
            if trigger_number < 0:
                trigger_number = get_trigger_number(raw_data_word, trigger_data_format=2)
            else:
                trigger_number = (0x7fffffffffff0000 & trigger_number) | get_trigger_number(raw_data_word, trigger_data_format=2)
            # Check validity of trigger number
            # Trigger number has to increase by 1
            if trigger_data_index >= 0:
                # Check if trigger number has increased by 1
                if last_trigger_number < 0:
                    n_missing_events = 0
                else:
                    n_missing_events = trigger_number - (last_trigger_number + 1)
                if n_missing_events != 0:
                    if n_missing_events > 0 and add_missing_events:
                        for i in range(n_missing_events):
                            # Increase index
                            trigger_data_index += 1
                            # extend trigger data array if neccessary
                            if trigger_data_index >= trigger_data.shape[0]:
                                trigger_data_tmp = np.zeros(shape=max(1, int(raw_data.shape[0] / 6)), dtype=trigger_data_dtype)
                                trigger_data = np.concatenate((trigger_data, trigger_data_tmp))
                            # Increase event number
                            event_number += 1
                            # Store trigger data
                            trigger_data[trigger_data_index]['event_number'] = event_number  # Timestamp of TLU word
                            trigger_data[trigger_data_index]['trigger_time_stamp'] = -1  # Timestamp of TLU word
                            trigger_data[trigger_data_index]['trigger_number'] = trigger_data[trigger_data_index - 1]['trigger_number'] + 1 + i
                            trigger_data[trigger_data_index]['trigger_status'] = NO_TRIGGER_WORD_ERROR  # Trigger status
                    else:
                        trigger_status |= TRIGGER_NUMBER_ERROR
            # Increase index
            trigger_data_index += 1
            # extend trigger data array if neccessary
            if trigger_data_index >= trigger_data.shape[0]:
                trigger_data_tmp = np.zeros(shape=max(1, int(raw_data.shape[0] / 6)), dtype=trigger_data_dtype)
                trigger_data = np.concatenate((trigger_data, trigger_data_tmp))
            # Increase event number
            event_number += 1
            # Store trigger data
            trigger_data[trigger_data_index]['event_number'] = event_number  # Timestamp of TLU word
            trigger_data[trigger_data_index]['trigger_number'] = trigger_number
            trigger_data[trigger_data_index]['trigger_time_stamp'] = trigger_timestamp  # Timestamp of TLU word
            trigger_data[trigger_data_index]['trigger_status'] = trigger_status  # Trigger status
        else:  # Raw data contains unknown word, neither M26 nor TLU word
            for tmp_plane_index, _ in enumerate(analyze_m26_header_ids):
                m26_data_loss[tmp_plane_index] = True

    # Set the status bits for priviously incomplete frames
    if build_all_events:
        for tmp_plane_index, tmp_plane_id in enumerate(analyze_m26_header_ids):
            index = telescope_data_index
            while index >= 0:
                if telescope_data[index]['plane'] == tmp_plane_id:
                    if telescope_data[index]['frame_id'] > last_completed_m26_frame_ids[tmp_plane_index]:
                        telescope_data[index]['frame_status'] |= DATA_ERROR
                    else:
                        break
                index -= 1

    return trigger_data, trigger_data_index, telescope_data, telescope_data_index, m26_frame_ids, m26_frame_length, m26_data_loss, m26_word_index, m26_timestamps, last_m26_timestamps, m26_n_words, m26_rows, m26_frame_status, last_completed_m26_frame_ids, event_number, trigger_number, trigger_timestamp

@njit(locals={'hits_index': numba.int64, 'curr_trigger_data_index': numba.int64, 'curr_telescope_data_index': numba.int64})
def _build_events(trigger_data, trigger_data_index, telescope_data, telescope_data_index, hits, hits_index, last_completed_m26_frame_ids, timing_offset, build_all_events, analyze_m26_header_ids, plane_id_to_index):
    ''' This function is builds events from the temporary trigger and telescope data arrays.

    Parameters:
    -----------
    TBD
    '''
    latest_trigger_data_index = -1
    finished_telescope_data_indices = -1 * np.ones(shape=len(analyze_m26_header_ids), dtype=np.int64)
    last_event_trigger_data_indices = -1 * np.ones(shape=len(analyze_m26_header_ids), dtype=np.int64)
    finished_event = np.ones(shape=len(analyze_m26_header_ids), dtype=np.bool_)
    curr_event_status = np.zeros(shape=len(analyze_m26_header_ids), dtype=np.uint32)

    curr_trigger_data_index = 0
    curr_hits_index = hits_index
    # adding hits
    while curr_trigger_data_index <= trigger_data_index and np.all(finished_event):
        trigger_event_number = trigger_data[curr_trigger_data_index]['event_number']
        trigger_number = trigger_data[curr_trigger_data_index]['trigger_number']
        trigger_timestamp = trigger_data[curr_trigger_data_index]['trigger_time_stamp']
        trigger_status = trigger_data[curr_trigger_data_index]['trigger_status']
        curr_telescope_data_index = np.min(finished_telescope_data_indices) + 1
        # Reset status
        for tmp_plane_index, _ in enumerate(analyze_m26_header_ids):
            finished_event[tmp_plane_index] = False
            curr_event_status[tmp_plane_index] = 0
        while curr_telescope_data_index <= telescope_data_index:
            curr_plane_id = telescope_data[curr_telescope_data_index]['plane']
            curr_plane_index = plane_id_to_index[curr_plane_id]
            curr_frame_id = telescope_data[curr_telescope_data_index]['frame_id']
            if not finished_event[curr_plane_index] and (build_all_events or curr_frame_id <= last_completed_m26_frame_ids[curr_plane_index]):
                hit_timestamp_start = telescope_data[curr_telescope_data_index]['time_stamp'] + telescope_data[curr_telescope_data_index]['row'] * ROW_UNIT_CYCLE - 2 * FRAME_UNIT_CYCLE - timing_offset
                hit_timestamp_stop = hit_timestamp_start + FRAME_UNIT_CYCLE + ROW_UNIT_CYCLE
                if hit_timestamp_start <= trigger_timestamp and trigger_timestamp < hit_timestamp_stop:
                    curr_hits_index += 1
                    # extend hits array if neccessary
                    if curr_hits_index >= hits.shape[0]:
                        hits_tmp = np.zeros(shape=telescope_data.shape[0], dtype=hits_dtype)
                        hits = np.concatenate((hits, hits_tmp))
                    # Adding hits to event
                    hits[curr_hits_index]['plane'] = curr_plane_id
                    hits[curr_hits_index]['event_number'] = trigger_event_number
                    hits[curr_hits_index]['trigger_number'] = trigger_number
                    hits[curr_hits_index]['trigger_time_stamp'] = trigger_timestamp
                    hits[curr_hits_index]['row_time_stamp'] = hit_timestamp_start
                    hits[curr_hits_index]['frame_id'] = curr_frame_id
                    hits[curr_hits_index]['column'] = telescope_data[curr_telescope_data_index]['column']
                    hits[curr_hits_index]['row'] = telescope_data[curr_telescope_data_index]['row']
                    hits[curr_hits_index]['event_status'] = 0
                    curr_event_status[curr_plane_index] |= telescope_data[curr_telescope_data_index]['frame_status'] | trigger_status
                elif hit_timestamp_start > trigger_timestamp:
                    # latest_trigger_data_indices[plane_id_to_index[telescope_data[curr_telescope_data_index]['plane']]] = curr_trigger_data_index
                    finished_event[plane_id_to_index[telescope_data[curr_telescope_data_index]['plane']]] = True
                    if np.all(finished_event):
                        latest_trigger_data_index = curr_trigger_data_index
                        for tmp_plane_index, _ in enumerate(analyze_m26_header_ids):
                            last_event_trigger_data_indices[tmp_plane_index] = finished_telescope_data_indices[tmp_plane_index]
                        hits_index = curr_hits_index
                        # Set event status for complete event
                        index = curr_hits_index
                        while trigger_event_number == hits[index]['event_number'] and index >= 0:
                            hits[index]['event_status'] = curr_event_status[plane_id_to_index[hits[index]['plane']]]
                            index -= 1
                        break
                else:  # trigger_timestamp >= hit_timestamp_stop
                    finished_telescope_data_indices[plane_id_to_index[telescope_data[curr_telescope_data_index]['plane']]] = curr_telescope_data_index
            curr_telescope_data_index += 1
        # special case
        if build_all_events:
            hits_index = curr_hits_index
            # Set event status for complete event
            index = curr_hits_index
            while index >= 0:
                if hits[index]['event_number'] == trigger_event_number:
                    hits[index]['event_status'] = curr_event_status[plane_id_to_index[hits[index]['plane']]]
                elif hits[index]['event_number'] < trigger_event_number:
                    break
                index -= 1
            for tmp_plane_index, _ in enumerate(analyze_m26_header_ids):
                finished_event[tmp_plane_index] = True
        curr_trigger_data_index += 1

    if build_all_events:
        telescope_data_start_index = telescope_data_index + 1
    else:
        telescope_data_start_index = np.min(last_event_trigger_data_indices) + 1
    telescope_data = telescope_data[telescope_data_start_index:]
    telescope_data_index -= telescope_data_start_index
    if build_all_events:
        trigger_data_start_index = trigger_data_index + 1
    else:
        trigger_data_start_index = latest_trigger_data_index + 1
    trigger_data = trigger_data[trigger_data_start_index:]
    trigger_data_index -= trigger_data_start_index

    return trigger_data, trigger_data_index, telescope_data, telescope_data_index, hits, hits_index


"""Calculate collision matrix of direct solution of LBTE."""
# Copyright (C) 2020 Atsushi Togo
# All rights reserved.
#
# This file is part of phono3py.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in
#   the documentation and/or other materials provided with the
#   distribution.
#
# * Neither the name of the phonopy project nor the names of its
#   contributors may be used to endorse or promote products derived
#   from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

import numpy as np
from phonopy.units import Kb, THzToEv

from phono3py.phonon3.imag_self_energy import ImagSelfEnergy
from phono3py.phonon3.interaction import Interaction

class CollisionMatrix(ImagSelfEnergy):
    """Collision matrix of direct solution of LBTE for one grid point.

    Main diagonal part (imag-self-energy) and
    the other part are separately stored.

    """

    def __init__(
        self,
        interaction: Interaction,
        rotations_cartesian=None,
        num_ir_grid_points=None,
        rot_grid_points=None,
        is_reducible_collision_matrix=False,
        log_level=0,
        lang="C",
    ):
        """Init method."""
        self._pp: Interaction
        self._is_collision_matrix: bool
        self._sigma = None
        self._frequency_points = None
        self._temperature = None
        self._grid_point = None
        self._lang = None
        self._imag_self_energy = None
        self._collision_matrix = None
        self._pp_strength = None
        self._frequencies = None
        self._triplets_at_q = None
        self._triplets_map_at_q = None
        self._weights_at_q = None
        self._band_indices = None
        self._unit_conversion = None
        self._cutoff_frequency = None
        self._g = None
        self._unit_conversion = None
        self._log_level = log_level

        super().__init__(interaction, lang=lang)

        self._is_reducible_collision_matrix = is_reducible_collision_matrix
        self._is_collision_matrix = True

        if not self._is_reducible_collision_matrix:
            self._num_ir_grid_points = num_ir_grid_points
            self._rot_grid_points = np.array(
                self._pp.bz_grid.bzg2grg[rot_grid_points], dtype="int_", order="C"
            )
            self._rotations_cartesian = rotations_cartesian

    def run(self):
        """Calculate collision matrix at a grid point."""
        if self._pp_strength is None:
            self.run_interaction()

        num_band0 = self._pp_strength.shape[1]
        num_band = self._pp_strength.shape[2]
        self._imag_self_energy = np.zeros(num_band0, dtype="double")

        if self._is_reducible_collision_matrix:
            num_mesh_points = np.prod(self._pp.mesh_numbers)
            self._collision_matrix = np.zeros(
                (num_band0, num_mesh_points, num_band), dtype="double"
            )
        else:
            self._collision_matrix = np.zeros(
                (num_band0, 3, self._num_ir_grid_points, num_band, 3), dtype="double"
            )
        self._run_with_band_indices()
        self._run_collision_matrix()

    def get_collision_matrix(self):
        """Return collision matrix at a grid point."""
        return self._collision_matrix

    def set_grid_point(self, grid_point=None):
        """Set a grid point and prepare for collision matrix calculation."""
        if grid_point is None:
            self._grid_point = None
        else:
            self._pp.set_grid_point(grid_point, store_triplets_map=True)
            self._pp_strength = None
            (
                self._triplets_at_q,
                self._weights_at_q,
                self._triplets_map_at_q,
                self._ir_map_at_q,
            ) = self._pp.get_triplets_at_q()
            self._grid_point = grid_point
            self._frequencies, self._eigenvectors, _ = self._pp.get_phonons()

    def _run_collision_matrix(self):
        if self._temperature > 0:
            if self._lang == "C":
                if self._is_reducible_collision_matrix:
                    self._run_c_reducible_collision_matrix()
                else:
                    self._run_c_collision_matrix()
            else:
                if self._is_reducible_collision_matrix:
                    self._run_py_reducible_collision_matrix()
                else:
                    self._run_py_collision_matrix()

    def _run_c_collision_matrix(self):
        import phono3py._phono3py as phono3c

        phono3c.collision_matrix(
            self._collision_matrix,
            self._pp_strength,
            self._frequencies,
            self._g,
            self._triplets_at_q,
            self._triplets_map_at_q,
            self._ir_map_at_q,
            self._rot_grid_points,  # in GRGrid
            self._rotations_cartesian,
            self._temperature,
            self._unit_conversion,
            self._cutoff_frequency,
        )

    def _run_c_reducible_collision_matrix(self):
        import phono3py._phono3py as phono3c

        phono3c.reducible_collision_matrix(
            self._collision_matrix,
            self._pp_strength,
            self._frequencies,
            self._g,
            self._triplets_at_q,
            self._triplets_map_at_q,
            self._ir_map_at_q,
            self._temperature,
            self._unit_conversion,
            self._cutoff_frequency,
        )

    def _run_py_collision_matrix(self):
        r"""Sum over rotations, and q-points and bands for third phonons.

        \Omega' = \sum_R' R' \Omega_{kp,R'k'p'}

        pp_strength.shape = (num_triplets, num_band0, num_band, num_band)

        """
        num_band0 = self._pp_strength.shape[1]
        num_band = self._pp_strength.shape[2]
        gp2tp, tp2s, swapped = self._get_gp2tp_map()
        for i in range(self._num_ir_grid_points):
            r_gps = self._rot_grid_points[i]
            for r, r_gp in zip(self._rotations_cartesian, r_gps):
                inv_sinh = self._get_inv_sinh(tp2s[r_gp])
                ti = gp2tp[r_gp]
                for j, k in np.ndindex((num_band0, num_band)):
                    if swapped[r_gp]:
                        collision = (
                            self._pp_strength[ti, j, :, k]
                            * inv_sinh
                            * self._g[2, ti, j, :, k]
                        ).sum()
                    else:
                        collision = (
                            self._pp_strength[ti, j, k]
                            * inv_sinh
                            * self._g[2, ti, j, k]
                        ).sum()
                    collision *= self._unit_conversion
                    self._collision_matrix[j, :, i, k, :] += collision * r

    def _run_py_reducible_collision_matrix(self):
        r"""Sum over q-points and bands of third phonons.

        This corresponds to the second term of right hand side of
        \Omega_{q0p0, q1p1} in Chaput's paper.

        pp_strength.shape = (num_triplets, num_band0, num_band, num_band)

        """
        num_mesh_points = np.prod(self._pp.mesh_numbers)
        num_band0 = self._pp_strength.shape[1]
        num_band = self._pp_strength.shape[2]
        gp2tp, tp2s, swapped = self._get_gp2tp_map()
        for gp1 in range(num_mesh_points):
            inv_sinh = self._get_inv_sinh(tp2s[gp1])
            ti = gp2tp[gp1]
            for j, k in np.ndindex((num_band0, num_band)):
                if swapped[gp1]:
                    collision = (
                        self._pp_strength[ti, j, :, k]
                        * inv_sinh
                        * self._g[2, ti, j, :, k]
                    ).sum()
                else:
                    collision = (
                        self._pp_strength[ti, j, k] * inv_sinh * self._g[2, ti, j, k]
                    ).sum()
                collision *= self._unit_conversion
                self._collision_matrix[j, gp1, k] += collision

    def _get_gp2tp_map(self):
        """Return mapping table from grid point index to triplet index.

        triplets_map_at_q contains index mapping of q1 in (q0, q1, q2) to
        independet q1 under q0+q1+q2=G with a fixed q0.

        Note
        ----
        map_q[gp1] <= gp1.:
            Symmetry relation of grid poi nts with a stabilizer q0.
        map_triplets[gp1] <= gp1 :
            map_q[gp1] == gp1 : map_q[gp2] if map_q[gp2] < gp1 otherwise gp1.
            map_q[gp1] != gp1 : map_triplets[map_q[gp1]]

        As a rule
        1. map_triplets[gp1] == gp1 : [gp0, gp1, gp2]
        2. map_triplets[gp1] != gp1 : [gp0, map_q[gp2], gp1'],
                                      map_triplets[gp1] == map_q[gp2]

        """
        map_triplets = self._triplets_map_at_q
        map_q = self._ir_map_at_q
        gp2tp = -np.ones(len(map_triplets), dtype="int_")
        tp2s = -np.ones(len(map_triplets), dtype="int_")
        swapped = np.zeros(len(map_triplets), dtype="bytes")
        num_tps = 0

        bzg2grg = self._pp.bz_grid.bzg2grg

        for gp1, tp_gp1 in enumerate(map_triplets):
            if map_q[gp1] == gp1:
                if gp1 == tp_gp1:
                    gp2tp[gp1] = num_tps
                    tp2s[gp1] = self._triplets_at_q[num_tps][2]
                    assert bzg2grg[self._triplets_at_q[num_tps][1]] == gp1
                    num_tps += 1
                else:  # q1 <--> q2 swap if swappable.
                    gp2tp[gp1] = gp2tp[tp_gp1]
                    tp2s[gp1] = self._triplets_at_q[gp2tp[gp1]][1]
                    swapped[gp1] = 1
                    assert map_q[bzg2grg[self._triplets_at_q[gp2tp[gp1]][2]]] == gp1
            else:  # q1 is not in ir-q1s.
                gp2tp[gp1] = gp2tp[map_q[gp1]]
                tp2s[gp1] = tp2s[map_q[gp1]]
                swapped[gp1] = swapped[map_q[gp1]]

            # Alternative implementation of tp2s
            # grg2bzg = self._pp.bz_grid.grg2bzg
            # addresses = self._pp.bz_grid.addresses
            # q0 = addresses[self._triplets_at_q[0][0]]
            # q1 = addresses[grg2bzg[gp1]]
            # q2 = -q0 - q1
            # gp2 = get_grid_point_from_address(q2, self._pp.bz_grid.D_diag)
            # tp2s[gp1] = self._pp.bz_grid.grg2bzg[gp2]

        return gp2tp, tp2s, swapped

    def _get_inv_sinh(self, gp):
        """Return sinh term for bands at a q-point."""
        freqs = self._frequencies[gp]
        sinh = np.where(
            freqs > self._cutoff_frequency,
            np.sinh(freqs * THzToEv / (2 * Kb * self._temperature)),
            -1.0,
        )
        inv_sinh = np.where(sinh > 0, 1.0 / sinh, 0)

        return inv_sinh


#! /usr/bin/env python
#S.rodney
# 2011.05.04
"""
Extrapolate the Hsiao SED and the non1a SEDs
down to 300 angstroms to allow the W filter to
reach out to z=2.5 smoothly in the k-correction
tables and model light curves.
"""

import os
from numpy import *
from pylab import *

try : sndataroot = os.environ['SNDATA_ROOT']
except KeyError:
    sndataroot = os.path.abspath( '.' )

MINWAVE = 300     # min wavelength for extrapolation (Angstroms)
MAXWAVE = 20000   # max wavelength for extrapolation (Angstroms)

def extendEverything( salt2dir = 'models/SALT2/SALT2.Guy10_UV2IR',
                      doCC=True, doIa=True ) :
    """ Run all the extrapolation functions, for both the SALT2 model
    components and all the non1a SEDs """

    if doCC :
        extendNon1a()

    if doIa :
        salt2dir = os.path.join( sndataroot, salt2dir )
        if not os.path.isdir( salt2dir ) :
            os.makedirs( salt2dir )

        # make the Hsiao07.extrap.dat template SED
        hsiao07 = os.path.join( sndataroot, 'snsed/Hsiao07.dat')
        hsiao07ext = os.path.join( sndataroot, 'snsed/Hsiao07.extrap.dat')
        extrapolatesed_linfit( hsiao07, hsiao07ext, Npt=10, Nsmooth=0, forceslope=True )

        extendSALT2_temp0()
        extendSALT2_temp1()
        extendSALT2_flatline()

def getsed( sedfile = os.path.join( sndataroot, 'snsed/Hsiao07.dat'), day='all'  ) :
    d,w,f = loadtxt( sedfile, unpack=True )

    #d = d.astype(int)
    days = unique( d )

    if day == 'all' :
        dlist = [ d[ where( d == day ) ] for day in days ]
        wlist = [ w[ where( d == day ) ] for day in days ]
        flist = [ f[ where( d == day ) ] for day in days ]
        return( dlist, wlist, flist )
    else :
        return( w[ where( d == day ) ], f[ where( d == day ) ] )

def plotsed( sedfile= os.path.join( sndataroot, 'snsed/Hsiao07.dat'),
             day='all', normalize=False, **kwarg):
    dlist,wlist,flist = getsed( sedfile )
    #days = unique( dlist )
    for i in range( len(wlist) ) :
        thisday = dlist[i][0]

        #defaults = { 'label':str(thisday) }
        #plotarg = dict( kwarg.items() + defaults.items() )
        if day!='all' :
            if abs(thisday-day)>0.6 : continue
        offset, normfactor = 0, 1
        if normalize :
            if normalize > 1 :
                normfactor = (flist[i][:-1] * (wlist[i][1:]-wlist[i][:-1])).sum()
            else :
                normfactor = flist[i].max()
                offset = thisday
        plot( wlist[i], flist[i]/normfactor + offset, **kwarg )
        # user_in=raw_input('%i : return to continue'%i)

def extrapolatesed( sedfile, newsedfile, minwave=MINWAVE, maxwave=MAXWAVE,
                    Bzptwave=900, Bextwave=2000, Bfrac=0.2, Brefwave=2800,
                    Nred = 10, Nsmooth=0, verbose=False ):
    """
    Set bluewave and bluefrac to define an extrapolation point
    blueward of the low-wavelength edge of the SED:
      Bextwave = wavelength of the extrapolation point
      Brefwave = reference wavelength for setting extrapolated flux
      Bfrac = fraction of the flux at Brefwave to set as the
          flux value at Bextwave
    The linear extrapolation  uses only this fixed point and
    the input SED point at Brefwave.

    On the red side (where the SED tail is more well behaved)
    we perform a linear fit to the last Nred points to define
    the extrapolation.  Set Nsmooth to a positive integer to set
    the window size for median smoothing of the input spectrum
    before the linear fitting.

    e.g.  Bzptwave, Bextwave, Bfrac, Brefwave = 900, 2000, 0.2, 2800
     fix the flux at 900 angstroms to be 0.
     fix the flux at 2000 angstroms to be 20% of the input SED's
     flux at 2800 angstrom, then define a piecewise linear extrapolation
     connecting the three anchor points at 900, 2000, and the minimum
     wavelength of the input SED.

    Note: For CCSN SEDs, we use
       Bzptwave=900, Bextwave=2000, Bfrac=0.2, Brefwave=2800,

    and for a SNIa SED, we might use
       Bzptwave=600, Bextwave=800, Bfrac=0.1, Brefwave=1000

    """
    from scipy import interpolate as scint
    from scipy import stats
    import shutil

    medsmooth = lambda f,N : array( [ median( f[max(0,i-N):min(len(f),max(0,i-N)+2*N)]) for i in range(len(f)) ] )

    dlist,wlist,flist = getsed( sedfile )  # these are 2-D lists of SEDs at different phases
    dlistnew, wlistnew, flistnew = [],[],[]

    fout = open( newsedfile, 'w' )
    for i in range( len(dlist) ) :
        d,w,f = dlist[i],wlist[i],flist[i]

        if minwave  < w[0] :
            wavestep = w[1] - w[0]
            irefpt = abs( w - Brefwave ).argmin()
            wrefpt = w[irefpt]
            frefpt = f[irefpt]
            wextpt = Bextwave
            fextpt = Bfrac * frefpt

            if wextpt < w[0] :
                # Blueward extrapolation from the minimum wavelength of the input
                # SED down to the first extrapolation point
                a = (f[0]-fextpt)/(w[0]-wextpt)  # slope of extrapolation line
                b = f[0] - a*w[0] # intercept
                wextBlue = arange( wextpt, w[0], wavestep )
                fextBlue = array( [ max( 0, a * wave + b ) for wave in wextBlue ] )
                w = append( wextBlue, w )
                f = append( fextBlue, f )
            if minwave < w[0] :
                # Blueward extrapolation to the min wavelength (fixed to zero flux at Bzptwave)
                wextpt = Bzptwave
                fextpt = 0
                a = (f[0]-fextpt)/(w[0]-wextpt)  # slope of extrapolation line
                b = f[0] - a*w[0] # intercept
                wextBlue = arange( minwave, w[0], wavestep )
                fextBlue = array( [ max( 0, a * wave + b ) for wave in wextBlue ] )
                w = append( wextBlue, w )
                f = append( fextBlue, f )

        if maxwave > w[-1] :
            # Redward linear extrapolation from last Nred points
            wavestep = w[-1] - w[-2]
            wN = w[-Nred:]

            if Nsmooth : fN = medsmooth( f, Nsmooth )[-Nred:]
            else : fN = f[-Nred:]
            (a,b,rval,pval,stderr)=stats.linregress(wN,fN)

            if a > 0 :
                # re-do redward  linear extrapolation using the peak
                # in order to guarantee a downward slope
                if verbose : print( "  re-fitting %s day %i to ensure downward slope to the red"%(sedfile, d[0]) )
                wN = [ w[ f.argmax() ], w[-1] ]
                fN = [ f.max(), f[-1],  ]
                (a,b,rval,pval,stderr)=stats.linregress(wN,fN)

            wextRed = arange( w[-1]+wavestep, maxwave, wavestep )
            fextRed = array( [ max( 0, a * wave + b ) for wave in wextRed ] )

            w = append( w, wextRed )
            f = append( f, fextRed )

        for i in range( len( w ) ) :
            print >> fout, "%7.2f  %10.2f  %12.7e"%( d[0], w[i], f[i] )
    fout.close()

    return( newsedfile )

def extrapolatesed_linfit(sedfile, newsedfile, minwave=MINWAVE, maxwave=MAXWAVE,
                          Npt=2, Nsmooth=0, forceslope=True, verbose=False ):
    """ use a linear fit of the first/last Npt  points on the SED
    to extrapolate to the blue and red.

    Set Nsmooth to a positive integer to set the window size for median
    smoothing of the input spectrum before the linear fitting.

    With forceslope=True, check if the linear fit returns a slope that would
    send the tail upward (i.e. a positive derivative on the red side, negative
    on the  blue side).  If so, then re-do the fit using just the first/last
    point and the peak flux point, in order to guarantee a downward sloping tail.
    """

    from scipy import interpolate as scint
    from scipy import stats
    import shutil

    medsmooth = lambda f,N : array( [ median( f[max(0,i-N):min(len(f),max(0,i-N)+2*N)]) for i in range(len(f)) ] )

    dlist,wlist,flist = getsed( sedfile )
    dlistnew, wlistnew, flistnew = [],[],[]

    fout = open( newsedfile, 'w' )
    for i in range( len(dlist) ) :
        d,w,f = dlist[i],wlist[i],flist[i]

        if Nsmooth : ffit = medsmooth( f, Nsmooth )
        else : ffit = f
        wavestep = w[1] - w[0]

        # blueward linear extrapolation from first N points
        wN = w[:Npt]
        fN = ffit[:Npt]
        (a,b,rval,pval,stderr)=stats.linregress(wN,fN)
        if forceslope and a < 0 :
            # re-do blueward linear extrapolation using the peak
            # in order to guarantee a downward slope
            print( "  re-fitting %s day %i to ensure downward slope to the blue"%(sedfile,d[0]))
            wN = [ w[0], w[ ffit.argmax() ] ]
            fN = [ f[0], ffit.max() ]
            (a,b,rval,pval,stderr)=stats.linregress(wN,fN)
        Nbluestep = len( arange( minwave, w[0], wavestep ) )
        wextBlue = sorted( [ w[0] -(i+1)*wavestep for i in range(Nbluestep) ] )
        fextBlue = array( [ max( 0, a * wave + b ) for wave in wextBlue ] )

        # redward linear extrapolation from last N points
        wN = w[-Npt:]
        fN = ffit[-Npt:]
        (a,b,rval,pval,stderr)=stats.linregress(wN,fN)
        if forceslope and a > 0 :
            # re-do redward  linear extrapolation using the peak
            # in order to guarantee a downward slope
            if verbose : print( "  re-fitting %s day %i to ensure downward slope to the red"%(sedfile, d[0]) )
            wN = [ w[ ffit.argmax() ], w[-1] ]
            fN = [ ffit.max(), ffit[-1],  ]
            (a,b,rval,pval,stderr)=stats.linregress(wN,fN)
        Nredstep = len( arange( w[-1], maxwave,  wavestep ) )
        wextRed =  sorted( [ w[-1] + (i+1)*wavestep for i in range(Nredstep) ] )
        fextRed = array( [ max( 0, a * wave + b ) for wave in wextRed ] )

        wnew = append( append( wextBlue, w ), wextRed )
        fnew = append( append( fextBlue, f ), fextRed )
        # dnew = zeros( len(wnew) ) + d[0]

        for i in range( len( wnew ) ) :
            print >> fout, "%7.2f  %10.2f  %12.7e"%( d[0], wnew[i], fnew[i] )
    fout.close()

    return( newsedfile )

def extrapolatesed_flatline(sedfile, newsedfile, minwave=MINWAVE, maxwave=MAXWAVE ):
    """ extrapolate to the red and the blue using a flatline. i.e. the
    extrapolated flux is fixed to the endpoint values
    """
    from scipy import interpolate as scint
    from scipy import stats
    import shutil

    dlist,wlist,flist = getsed( sedfile )
    dlistnew, wlistnew, flistnew = [],[],[]

    fout = open( newsedfile, 'w' )
    for i in range( len(dlist) ) :
        d,w,f = dlist[i],wlist[i],flist[i]

        wavestep = w[1] - w[0]
        # blueward flatline extrapolation from first point
        Nbluestep = len( arange( minwave, w[0], wavestep ) )
        wextBlue = sorted( [ w[0] -(i+1)*wavestep for i in range(Nbluestep) ] )
        fextBlue = array( [ f[0] for wave in wextBlue ] )

        # redward flatline extrapolation from last point
        Nredstep = len( arange( w[-1], maxwave,  wavestep ) )
        wextRed =  sorted( [ w[-1] + (i+1)*wavestep for i in range(Nredstep) ] )
        fextRed = array( [ f[-1]  for wave in wextRed ] )

        wnew = append( append( wextBlue, w ), wextRed )
        fnew = append( append( fextBlue, f ), fextRed )
        # dnew = zeros( len(wnew) ) + d[0]

        for i in range( len( wnew ) ) :
            print >> fout, "%5.1f  %10i  %12.7e"%( d[0], wnew[i], fnew[i] )
    fout.close()

    return( newsedfile )

def extendNon1a(origdir = "snsed/non1a.ORIG", verbose=True ):
    """ Extrapolate each of the Non-Ia template sed files,
    preserving a pristine copy in $SNDATA_ROOT/snsed/non1a.ORIG
    """
    import glob
    import shutil

    origdir = os.path.join(sndataroot, origdir)
    non1adir = os.path.join(sndataroot, "snsed/non1a")
    if not os.path.isdir( origdir ) :
        os.rename( non1adir, origdir )
    if not os.path.isdir( non1adir ) :
        os.mkdir( non1adir )

    otherstufflist = glob.glob("%s/*.DAT"%(origdir)) + glob.glob("%s/*.LIST"%(origdir)) + glob.glob("%s/*.INPUT"%(origdir))
    for otherfile in otherstufflist :
        newfile =  os.path.join( non1adir, os.path.basename( otherfile ) )
        shutil.copy( otherfile, newfile )

    sedlist = glob.glob("%s/*.SED"%origdir)
    for sedfile in sedlist :
        newsedfile =  os.path.join( non1adir, os.path.basename( sedfile ) )

        if 'SDSS-012842' in sedfile or 'SDSS-013449' in sedfile :
            # These are already super-smooth blackbodies for the IIn templates
            # 2-pt linear extrapolation is best.
            Npt = 2
            Nsmooth=0
            print("EXTRAPOLATING %s\n    ==> %s"%(sedfile, newsedfile) )
            extrapolatesed_linfit(sedfile, newsedfile, minwave=MINWAVE, maxwave=MAXWAVE,
                                  Npt=Npt, Nsmooth=Nsmooth, forceslope=True  )
            print("     Done with %s.\a\a\a"%newsedfile)

        else :
            # for all other templates, we use a fixed blue anchor point
            # and on the red side we fit a lot of points so we don't
            # get thrown around by sharp spectral features
            Nred = 100
            Nsmooth = 20

            print("EXTRAPOLATING %s\n    ==> %s"%(sedfile, newsedfile) )
            extrapolatesed( sedfile, newsedfile, minwave=MINWAVE, maxwave=MAXWAVE,
                            Bextwave=2000, Bfrac=0.2, Brefwave=2800,
                            Nred = 100, Nsmooth=20, verbose=verbose )
            print("     Done with %s.\a\a\a"%newsedfile)

def extendSALT2_temp0( salt2dir = 'models/SALT2/SALT2.Guy10_UV2IR',
                       salt2srcdir = 'models/SALT2/SALT2.Guy10_LAMOPEN',
                       tailsedfile = 'snsed/Hsiao07.extrap.dat',
                       wjoinblue = 2800, wjoinred = 8500 ,
                       wmin = MINWAVE, wmax = MAXWAVE ):
    """ extend the salt2 Template_0 model component
    by adopting the UV and IR tails from another SED model.
    The default is to use SR's extrapolated modification
    of the Hsiao 2007 sed model, scaled and joined at the
    wjoin wavelengths, and extrapolated out to wmin and wmax.
    """
    import shutil
    sndataroot = os.environ['SNDATA_ROOT']

    salt2dir = os.path.join( sndataroot, salt2dir )
    salt2srcdir = os.path.join( sndataroot, salt2srcdir )

    temp0fileIN = os.path.join( salt2srcdir, 'salt2_template_0.dat' )
    temp0fileOUT = os.path.join( salt2dir, 'salt2_template_0.dat' )
    temp0dat = getsed( sedfile=temp0fileIN )

    tailsedfile = os.path.join( sndataroot, tailsedfile )

    taildat = getsed( sedfile=tailsedfile )

    dt,wt,ft = loadtxt( tailsedfile, unpack=True )
    taildays = unique( dt )

    # build up modified template from day -20 to +50
    outlines = []
    for i in range( 71 ) :
        thisday = i - 20

        # get the tail SED for this day
        it = where( taildays == thisday )[0]
        dt = taildat[0][it]
        wt = taildat[1][it]
        ft = taildat[2][it]

        #if thisday > 50 :
        #    d0new = dt
        #    w0new = wt
        #    f0new = ft * (bluescale+redscale)/2.
        #else :

        # get the SALT2 template SED for this day
        d0 = temp0dat[0][i]
        w0 = temp0dat[1][i]
        f0 = temp0dat[2][i]
        print( 'splicing tail onto template for day : %i'%thisday )

        i0blue = argmin(  abs(w0-wjoinblue) )
        itblue = argmin( abs( wt-wjoinblue))

        i0red = argmin(  abs(w0-wjoinred) )
        itred = argmin( abs( wt-wjoinred))

        itmin = argmin( abs( wt-wmin))
        itmax = argmin( abs( wt-wmax))

        bluescale = f0[i0blue]/ft[itblue]
        redscale = f0[i0red]/ft[itred]

        d0new = dt.tolist()[itmin:itblue] + d0.tolist()[i0blue:i0red-1] + dt.tolist()[itred:itmax+1]
        w0new = wt.tolist()[itmin:itblue] + w0.tolist()[i0blue:i0red-1] + wt.tolist()[itred:itmax+1]
        f0new = (bluescale*ft).tolist()[itmin:itblue] + f0.tolist()[i0blue:i0red-1] + (redscale*ft).tolist()[itred:itmax+1]

        # plot it
        clf()
        plot( w0, f0, ls='-',color='b', lw=1, label='Input SALT2 Model')
        plot( wt, (bluescale+redscale)/2. * ft, ls=':',color='r', lw=1, label='SED model')
        plot( w0new, f0new, ls='--',color='k', lw=2, label='Extrapolated SALT2 model')
        legend()
        draw()
        #raw_input('return to continue')

        # append to the list of output data lines
        for j in range( len( d0new ) ) :
            outlines.append( "%6.2f    %12i  %12.7e\n"%(
                    d0new[j], w0new[j], f0new[j] ) )

    # write it out to the new template sed .dat file
    fout = open( temp0fileOUT, 'w' )
    fout.writelines( outlines )
    fout.close()

def extendSALT2_temp1( salt2dir = 'models/SALT2/SALT2.Guy10_UV2IR',
                       salt2srcdir = 'models/SALT2/SALT2.Guy10_LAMOPEN',
                       # wjoinblue = 2000, wjoinred = 8500 ,
                       wjoinblue = None, wjoinred = None,
                       wmin = MINWAVE, wmax = MAXWAVE,
                       wstep = 10 ):
    """ extend the salt2 Template_1 model component
    with a flat line at 0 to the blue and to the red.
    If join wavelengths are not provided, uses the
    end points of the input spectrum
    """
    import shutil
    sndataroot = os.environ['SNDATA_ROOT']

    salt2dir = os.path.join( sndataroot, salt2dir )
    salt2srcdir = os.path.join( sndataroot, salt2srcdir )

    temp1fileIN = os.path.join( salt2srcdir, 'salt2_template_1.dat' )
    temp1fileOUT = os.path.join( salt2dir, 'salt2_template_1.dat' )
    temp1dat = getsed( sedfile=temp1fileIN )

    # build up modified template from day -20 to +50
    outlines = []
    for i in range( 71 ) :
        thisday = i - 20

        # get the SALT2 template SED for this day
        d1 = temp1dat[0][i]
        w1 = temp1dat[1][i]
        f1 = temp1dat[2][i]
        print( 'extrapolating with flatline onto template for day : %i'%thisday )

        if wjoinblue==None : wjoinblue=w1[0]
        if wjoinred==None : wjoinred=w1[-1]

        i1blue = argmin(  abs(w1-wjoinblue) )
        i1red = argmin(  abs(w1-wjoinred) )

        Nblue = int((wjoinblue-wmin )/wstep + 1)
        Nred = int((wmax -wjoinred )/wstep + 1)

        d1new =  (ones(Nblue)*thisday).tolist() + d1.tolist()[i1blue+1:i1red-1] + (ones(Nred)*thisday).tolist()
        w1new = arange(wmin,wmin+Nblue*wstep,wstep).tolist() + w1.tolist()[i1blue+1:i1red-1] + arange(wjoinred,wjoinred+Nred*wstep,wstep).tolist()
        f1new = zeros(Nblue).tolist() + f1.tolist()[i1blue+1:i1red-1] + zeros(Nred).tolist()

        # plot it
        clf()
        plot( w1, f1, ls='-',color='r', lw=1)
        plot( w1new, f1new, ls='--',color='k', lw=2)
        draw()
        #raw_input('return to continue')

        # append to the list of output data lines
        for j in range( len( d1new ) ) :
            outlines.append( "%6.2f    %12i  %12.7e\n"%(
                    d1new[j], w1new[j], f1new[j] ) )

    # write it out to the new template sed .dat file
    fout = open( temp1fileOUT, 'w' )
    fout.writelines( outlines )
    fout.close()

def extendSALT2_flatline( salt2dir = 'models/SALT2/SALT2.Guy10_UV2IR',
                          salt2srcdir = 'models/SALT2/SALT2.Guy10_LAMOPEN',
                          wjoinblue = 2000, wjoinred = 8500 ,
                          wmin = MINWAVE, wmax = MAXWAVE,
                          wstep = 10, showplots=False ):
    """ extrapolate the *lc* and *spec* .dat files for SALT2
    using a flatline to the blue and red """

    sndataroot = os.environ['SNDATA_ROOT']
    salt2dir = os.path.join( sndataroot, salt2dir )
    salt2srcdir = os.path.join( sndataroot, salt2srcdir )

    filelist = ['salt2_lc_dispersion_scaling.dat',
                'salt2_lc_relative_covariance_01.dat',
                'salt2_lc_relative_variance_0.dat',
                'salt2_lc_relative_variance_1.dat',
                'salt2_spec_covariance_01.dat',
                'salt2_spec_variance_0.dat',
                'salt2_spec_variance_1.dat']

    #for filename in  ['salt2_lc_dispersion_scaling.dat']:
    #for filename in  ['salt2_lc_relative_covariance_01.dat']:
    for filename in filelist :
        infile = os.path.join( salt2srcdir, filename )
        outfile = os.path.join( salt2dir, filename )

        newsedfile = extrapolatesed_flatline( infile, outfile, minwave=wmin, maxwave=wmax,
                                              wjoinblue=wjoinblue, wjoinred=wjoinred )

        # plot it
        if showplots:
            #for d in range(-20,50) :
            for d in [-10,-5,0,5,10,15,20,25,30,35,40,45,50] :
                clf()
                plotsed( infile, day=d, ls='-',color='r', lw=1)
                plotsed( outfile,day=d, ls='--',color='k', lw=2)
                print( '%s : day %i'%(filename,d) )
                draw()
                # raw_input('%s : day %i.  return to continue'%(filename,d))

def plotAll():
    """
    read in and plot all the non1a template SEDs at peak,
    in three panels, with the Ia template for comparison,
    """
    ioff()
    ax1 = subplot( 131 )
    # plot the type II-L in purple
    plotsed( os.path.join(sndataroot,'snsed/non1a/Nugent+Scolnic_IIL.SED'), day=0, normalize=2, color='b', lw=0.5, label='II-L' )
    # plot the type IIn's in orange
    for sed in ['SDSS-012842','SDSS-013449'] :
        plotsed( os.path.join(sndataroot,'snsed/non1a/%s.SED'%sed), day=0.5, normalize=2, color='b', lw=0.5, label=sed )
    # Ia in red
    plotsed( os.path.join(sndataroot,'snsed/Hsiao07.extrap.dat'), day=0, normalize=2, color='r', label='Ia (Hsiao)', lw=2 )
    plotsed( os.path.join(sndataroot,'snsed/Hsiao07.dat'), day=0, normalize=2, color='k', label='Ia (Hsiao)', lw=1, ls=':' )
    title('Type II-L and IIn Templates')
    #legend()
    draw()

    subplot( 132, sharex=ax1, sharey=ax1 )
    # plot the type IIP's in blue
    IIpsedlist = [ 'SDSS-000018','SDSS-003818','SDSS-013376','SDSS-014450','SDSS-014599','SDSS-015031','SDSS-015320',
                   'SDSS-015339','SDSS-017564','SDSS-017862','SDSS-018109','SDSS-018297','SDSS-018408','SDSS-018441',
                   'SDSS-018457','SDSS-018590','SDSS-018596','SDSS-018700','SDSS-018713','SDSS-018734','SDSS-018793',
                   'SDSS-018834','SDSS-018892','SDSS-020038' ]
    for sed in IIpsedlist :
        plotsed( os.path.join(sndataroot,'snsed/non1a/%s.SED'%sed), day=0.5, normalize=2, color='b', lw=0.5, label=sed )
    # Ia in red
    plotsed( os.path.join(sndataroot,'snsed/Hsiao07.extrap.dat'), day=0, normalize=2, color='r', label='Ia (Hsiao)', lw=2 )
    plotsed( os.path.join(sndataroot,'snsed/Hsiao07.dat'), day=0, normalize=2, color='k', label='Ia (Hsiao)', lw=1, ls=':' )
    #legend()
    title('Type II-P Templates')
    draw()

    subplot( 133, sharex=ax1, sharey=ax1  )
    Ibclist =  [ 'CSP-2004gv','CSP-2006ep','CSP-2007Y','SDSS-000020','SDSS-002744','SDSS-014492',
                 'SDSS-019323','SNLS-04D1la','SNLS-04D4jv','CSP-2004fe','CSP-2004gq','SDSS-004012',
                 'SDSS-013195','SDSS-014475','SDSS-015475','SDSS-017548' ]
    for sed in Ibclist :
        plotsed( os.path.join(sndataroot,'snsed/non1a/%s.SED'%sed), day=0.5, normalize=2, color='g', lw=0.5, label=sed )
    # Ia in red
    plotsed( os.path.join(sndataroot,'snsed/Hsiao07.extrap.dat'), day=0, normalize=2, color='r', label='Ia (Hsiao)', lw=2 )
    plotsed( os.path.join(sndataroot,'snsed/Hsiao07.dat'), day=0, normalize=2, color='k', label='Ia (Hsiao)', lw=1, ls=':' )
    title('Type Ib and Ic Templates', color='g')
    #legend()
    ion()

    suptitle( 'Comparing SNANA CC and Ia templates at peak' )
    draw()

def extratest( sed, **kwarg ) :
    sedorig = '/usr/local/SNDATA_ROOT/snsed/non1a.ORIG/'+sed+'.SED'
    sednew  = '/usr/local/SNDATA_ROOT/snsed/non1a/'+sed+'test.SED'
    extrapolatesed_linfit( sedorig, sednew, **kwarg )
    clf()
    plotsed( sednew, day=0.5, color='g', ls='-', lw=1, normalize=2, label='new'  )
    plotsed( sedorig, day=0.5, color='b', ls='-', lw=3, normalize=2, label='orig'  )

    plotsed( '/usr/local/SNDATA_ROOT/snsed/Hsiao07.extrap.dat', day=0, normalize=2, color='r', label='Ia (Hsiao)', lw=2 )

    ax = gca()
    text( 0.1,0.95, sed, ha='left',va='top', transform=ax.transAxes )
    legend( )


import os
import numpy as np
from .decoder import AmbiDecoder
from .hrir import CIPIC_HRIR
from .position import Position, PositionalSource, MovingSource
from scipy.ndimage.interpolation import shift
#from tdesigns import get_tDesign

C = 343.     # Speed of sound [m/s]

class VirtualStereoMic(object):
    def __init__(self, radius=0.1):
        self.radius = radius
        self.lmic_pos = Position(0, radius, 0, 'cartesian')
        self.rmic_pos = Position(0, -radius, 0, 'cartesian')

    def binauralize(self, sources):
        if isinstance(sources, PositionalSource):
            sources = [sources]

        l_signal, r_signal = 0, 0.
        for src in sources:
            l_dist = np.sqrt(((src.position.coords('cartesian') - self.lmic_pos.coords('cartesian'))**2).sum())
            r_dist = np.sqrt(((src.position.coords('cartesian') - self.rmic_pos.coords('cartesian'))**2).sum())

            # Time delay
            l_delay, r_delay = int(l_dist / C * src.sample_rate), int(r_dist / C * src.sample_rate)

            # Attenuation is frequency dependent, but lets simplify.
            l_attn, r_attn = 1 / (1. + l_dist), 1 / (1. + r_dist)

            l_signal += l_attn * shift(src.signal, l_delay, cval=0.) / len(sources)
            r_signal += r_attn * shift(src.signal, r_delay, cval=0.) / len(sources)

        return np.stack((l_signal, r_signal), axis=1)

    def binauralize_frame(self, sources, output, frame_no):
        if isinstance(sources, PositionalSource):
            sources = [sources]

        for src in sources:
            l_dist = np.sqrt(((src.position.coords('cartesian') - self.lmic_pos.coords('cartesian'))**2).sum())
            r_dist = np.sqrt(((src.position.coords('cartesian') - self.rmic_pos.coords('cartesian'))**2).sum())

            # Time delay
            l_delay, r_delay = int(l_dist / C * src.sample_rate), int(r_dist / C * src.sample_rate)

            # Attenuation is frequency dependent, but lets simplify.
            l_attn, r_attn = 1 / (1. + l_dist), 1 / (1. + r_dist)

            if frame_no-l_delay >= 0:
                output[frame_no, 0] += l_attn * src.signal[frame_no-l_delay] / len(sources)
            if frame_no-r_delay >= 0:
                output[frame_no, 1] += r_attn * src.signal[frame_no-r_delay] / len(sources)

class Convolvotron(object):
    def __init__(self, cipic_dir):
        assert os.path.exists(cipic_dir)
        self.hrir_db = CIPIC_HRIR(cipic_dir)

    def binauralize(self, sources):
        if isinstance(sources, PositionalSource):
            sources = [sources]
        num_frames = max([src.signal.shape[0] for src in sources])
        stereo = np.zeros((num_frames, 2))
        for src in sources:
            left_hrir, right_hrir = self.hrir_db.get_closest(src.position)[1:]
            left_signal = np.convolve(src.signal, np.flip(left_hrir, axis=0), 'valid')
            right_signal = np.convolve(src.signal, np.flip(right_hrir, axis=0), 'valid')

            n_valid, i_start = left_signal.shape[0], left_hrir.shape[0] - 1
            stereo[i_start:(i_start + n_valid), 0] += left_signal
            stereo[i_start:(i_start + n_valid), 1] += right_signal

        return stereo

    def binauralize_frame(self, sources, output, frame_no):
        if isinstance(sources, PositionalSource):
            sources = [sources]

        for src in sources:
            left_hrir, right_hrir = self.hrir_db.get_closest(src.position)[1:]

            i_start = frame_no - left_hrir.size + 1 if frame_no >= left_hrir.size else 0
            i_end = frame_no + 1
            i_range = i_end - i_start

            output[frame_no, 0] = (src.signal[i_start:i_end] * left_hrir[-i_range:]).sum()
            output[frame_no, 1] = (src.signal[i_start:i_end] * right_hrir[-i_range:]).sum()

class SourceBinauralizer(object):
    def __init__(self, use_hrtfs=True, cipic_dir=None):
        self.use_hrts = use_hrtfs
        if use_hrtfs:
            self.convolvotron = Convolvotron(cipic_dir)
        else:
            self.stereo_mic = VirtualStereoMic()

    def binauralize(self, sources):
        if isinstance(sources, PositionalSource):
            sources = [sources]
        assert isinstance(sources, list) and all([isinstance(src, PositionalSource) for src in sources])
        assert all([src.sample_rate == sources[0].sample_rate for src in sources])

        if self.use_hrts:
            return self.convolvotron.binauralize(sources)
        else:
            return self.stereo_mic.binauralize(sources)

    def binauralize_frame(self, sources, output, frame_no):
        if isinstance(sources, PositionalSource):
            sources = [sources]
        assert isinstance(sources, list) and all([isinstance(src, PositionalSource) for src in sources])
        assert all([src.sample_rate == sources[0].sample_rate for src in sources])

        if self.use_hrts:
            return self.convolvotron.binauralize_frame(sources, output, frame_no)
        else:
            return self.stereo_mic.binauralize_frame(sources, output, frame_no)

class AmbisonicBinauralizer(object):
    def __init__(self, ambi_format, method='projection', use_hrtfs=False, cipic_dir=None):
        self.source_bin = SourceBinauralizer(cipic_dir=cipic_dir, use_hrtfs=use_hrtfs)
        self.fmt = ambi_format
        self.method = method

        # Initialize speakers
        if self.method == 'pseudoinv':
            self.speaker_pos = map(lambda x: Position(x[0], x[1], x[2], 'cartesian'), get_tDesign(self.fmt.order))
            map(lambda p: p.set_radius(self.fmt.radius), self.speaker_pos)
            # speakers_phi = (2. * np.arange(2*self.fmt.num_channels) / float(2*self.fmt.num_channels) - 1.) * np.pi
            # self.speaker_pos = map(lambda x: Position(x, 0, self.fmt.radius, 'polar'), speakers_phi)
        elif self.method == 'projection':
            speakers_phi = (2. * np.arange(2*self.fmt.num_channels) / float(2*self.fmt.num_channels) - 1.) * np.pi
            # self.speaker_pos = map(lambda x: Position(x, 0, self.fmt.radius, 'polar'), speakers_phi)
            self.speaker_pos = [Position(x, 0, self.fmt.radius, 'polar') for x in speakers_phi]
        else:
            raise ValueError('Unknown decoding method. Options: projection and pseudoinv')
        self.n_speakers = len(self.speaker_pos)
        self.ambi_decoder = AmbiDecoder(self.speaker_pos, self.fmt, method=self.method)

    def binauralize(self, ambi):
        # Decode ambisonics into speakers
        speakers = self.ambi_decoder.decode(ambi)

        # Binauralize speaker as if they were point sources
        sources = [PositionalSource(speakers[:, i], self.speaker_pos[i], self.fmt.sample_rate) for i in range(self.n_speakers)]
        stereo = self.source_bin.binauralize(sources)

        return stereo

class DirectAmbisonicBinauralizer(object):
    def __init__(self, ambi_format, method='projection'):
        self.fmt = ambi_format
        self.method = method

        # Initialize ear position
        self.ear_pos = [Position(0, 0.1, 0, 'cartesian'), Position(0, -0.1, 0, 'cartesian')]
        self.ambi_decoder = AmbiDecoder(self.ear_pos, self.fmt, method=self.method)

    def binauralize(self, ambi):
        return self.ambi_decoder.decode(ambi)

def test_convolvotron():
    from pyutils.iolib.audio import load_wav, save_wav
    convolvotron = Convolvotron('hrtfs/cipic_subj3')
    mono, rate = load_wav('wav_test/piano.wav')
    mono = mono[:, 0]

    positions = [[float(num) for num in l.strip().split()] for l in open('wav_test/piano_stat_position.txt', 'r')]
    positions = [Position(p[0], p[1], p[2], 'polar') for p in positions]
    source = PositionalSource(mono, positions[0], rate)

    stereo = convolvotron.binauralize([source])
    save_wav('/tmp/output.wav', stereo, rate)
    os.system('play /tmp/output.wav')
    os.remove('/tmp/output.wav')

    positions = [[float(num) for num in l.strip().split()] for l in open('wav_test/piano_mov_position.txt', 'r')]
    positions = [Position(p[0], p[1], p[2], 'polar') for p in positions]
    source = MovingSource(mono, positions, rate)

    stereo = np.zeros((mono.shape[0], 2))
    while source.tic():
        convolvotron.binauralize_frame([source], stereo, source.cur_idx)
    save_wav('/tmp/output.wav', stereo, rate)
    os.system('play /tmp/output.wav')
    os.remove('/tmp/output.wav')

def test_virtual_mic():
    from pyutils.iolib.audio import load_wav, save_wav
    mic = VirtualStereoMic()
    mono, rate = load_wav('wav_test/piano.wav')
    mono = mono[:, 0]

    positions = [[float(num) for num in l.strip().split()] for l in open('wav_test/piano_stat_position.txt', 'r')]
    positions = [Position(p[0], p[1], p[2], 'polar') for p in positions]
    source = PositionalSource(mono, positions[0], rate)

    stereo = mic.binauralize([source])
    save_wav('/tmp/output.wav', stereo, rate)
    os.system('play /tmp/output.wav')
    os.remove('/tmp/output.wav')

    positions = [[float(num) for num in l.strip().split()] for l in open('wav_test/piano_mov_position.txt', 'r')]
    positions = [Position(p[0], p[1], p[2], 'polar') for p in positions]
    source = MovingSource(mono, positions, rate)

    stereo = np.zeros((mono.shape[0], 2))
    while source.tic():
        mic.binauralize_frame([source], stereo, source.cur_idx)
    save_wav('/tmp/output.wav', stereo, rate)
    os.system('play /tmp/output.wav')
    os.remove('/tmp/output.wav')

def test_source_binauralizer():
    from pyutils.iolib.audio import load_wav, save_wav
    from pyutils.iolib.position import read_position_file

    # binauralizer = SourceBinauralizer(use_hrtfs=True, cipic_dir='hrtfs/cipic_subj3')
    binauralizer = SourceBinauralizer(use_hrtfs=False)

    # Static source
    sample = 'wav_test/gen_synthetic-S1'
    positions, wav_fns, _, sample_ids = read_position_file(sample+'-position.txt')
    mono, rate = load_wav(wav_fns[sample_ids[0]])
    source = PositionalSource(mono[:, 0], positions[sample_ids[0]][0], rate)
    stereo = binauralizer.binauralize([source])

    save_wav('/tmp/output.wav', stereo / np.abs(stereo).max(), rate)
    os.system('play /tmp/output.wav')
    os.remove('/tmp/output.wav')

    # Moving source
    sample = 'wav_test/gen_synthetic-M1'
    positions, wav_fns, _, sample_ids = read_position_file(sample+'-position.txt')
    mono, rate = load_wav(wav_fns[sample_ids[0]])
    source = MovingSource(mono[:, 0], positions[sample_ids[0]], rate)
    stereo = np.zeros((mono.shape[0], 2))
    while source.tic():
        binauralizer.binauralize_frame([source], stereo, source.cur_idx)

    save_wav('/tmp/output.wav', stereo / np.abs(stereo).max(), rate)
    os.system('play /tmp/output.wav')
    os.remove('/tmp/output.wav')

def test_ambisonics_binauralizer():
    from pyutils.iolib.audio import load_wav, save_wav
    from pyutils.ambisonics.common import AmbiFormat

    sample = 'wav_test/gen_synthetic-S1'
    ambi, rate = load_wav(sample+'-ambix.wav')

    fmt = AmbiFormat(1, rate)
    binauralizer = DirectAmbisonicBinauralizer(fmt, method='pseudoinv')
    # binauralizer = AmbisonicBinauralizer(fmt, method='projection', use_hrtfs=True, cipic_dir='hrtfs/cipic_subj3')

    stereo = binauralizer.binauralize(ambi)

    save_wav('/tmp/output.wav', stereo / np.abs(stereo).max(), rate)
    os.system('play /tmp/output.wav')
    os.remove('/tmp/output.wav')

    sample = 'wav_test/gen_synthetic-M1'
    ambi, rate = load_wav(sample+'-ambix.wav')
    stereo = binauralizer.binauralize(ambi)

    save_wav('/tmp/output.wav', stereo / np.abs(stereo).max(), rate)
    os.system('play /tmp/output.wav')
    os.remove('/tmp/output.wav')

if __name__ == '__main__':
    # test_ambisonics_binauralizer()
    test_source_binauralizer()
    # test_convolvotron()
    # test_virtual_mic()