# Import the new package made specifically for this analysis
from src.utils.paths import REPORTS_DIR

import src.features.chunk_onset_dev as cod

import BirdSongToolbox.free_epoch_tools as fet
from BirdSongToolbox.import_data import ImportData
from BirdSongToolbox.file_utility_functions import _handle_data_path, _save_pckl_data

import src.analysis.hilbert_based_pipeline as hbp
from src.analysis.ml_pipeline_utilities import all_bad_channels, all_drop_temps, all_label_instructions
from src.analysis.context_utility import birds_context_obj
from src.analysis.chunk_feature_dropping_pearson import best_bin_width, best_offset
import src.analysis.chunk_when_analysis_naive as cwan

import numpy as np
import scipy
from sklearn.model_selection import KFold
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

from matplotlib.backends.backend_pdf import PdfPages

# Location to save Data
branch_analysis_path = "/home/debrown/branch_analysis_path"

# Data Source for Script
onset_detection_path = '/home/debrown/onset_detection_results'

from BirdSongToolbox.file_utility_functions import _load_pckl_data

all_label_instruct = {'z007': [5],
                      'z020': [4],
                      'z017': [6, 7]
                      }

def naive_branch_analysis(absolute_relative_starts, motif_ledger, all_best_chunk_events, motif_time_series,
                          time_buffer, label_instruct):
    """ Algorithm to run a naive onset detection/predicition analysis.

    Parameters
    ----------
    absolute_relative_starts : list | [labels] -> [instances]
        list of all of the relative start times for each label in the order of labels_used
    motif_ledger : dict | { label_num : [Motif_identity]}
        A ledger of which motif a syllable instance occurred during. It is in order the syllables occurred assuming they
        started with the first syllable
    all_best_chunk_events : list | [Freq]->[labels]->(Instances, 1, Channels, Samples)
        Data clipped into the trials that can be used to make the Pearson Templates
    motif_time_series : list | [Frequency]->(Instances, Channels, Time-Steps, Bin Width)
        list of the times overlapping bins for each frequency to do the mini-onset detection
    time_buffer : int
        #Window around true start time

    label_instruct : list
        list of the syllable labels to use for the Particular bird

    Returns
    -------
    [syllables]->[Freqs]->[instances]

    elections: | [syllables]->[Freqs]->(instances, samples)
    """

    super_all_elections = []

    for sel_syll in label_instruct:

        syll_all_elections = []
        for sel_freq in range(len(all_best_chunk_events)):
            sel_label = sel_syll - 1

            expected_time = int(np.mean(absolute_relative_starts[sel_label - 1] / 30))
            expected_before = expected_time - time_buffer
            expected_after = expected_time + time_buffer
            print(expected_before, 'to', expected_after)

            syll_ex = np.arange(len(motif_ledger[sel_syll]))  # Index fo all of the Syllable examples
            branch_ex = [x for x in motif_ledger[2] if x not in motif_ledger[sel_syll]]  # Index of Motifs w/o the Syll

            # 1.) Get the Templates
            # 1.1) Sub-select Templates using Template Index
            sel_for_template = all_best_chunk_events[sel_freq][sel_label][syll_ex]
            sel_template = np.squeeze(np.mean(sel_for_template, axis=0))  # 1.2) Take the Mean (Chan, Samples)

            # 3.) Get the Time-Series [syllable](Instances, channels, samples, bin_width)
            # 3.1) Sub-select Motif Time Series using Test Index
            sel_time_series = motif_time_series[sel_freq][branch_ex]

            # 4.) Run the Test
            # 4.1) Get the Pearson Coefficients
            when_results = cod.get_freq_pearson_coefficients(freq_template_data=sel_template,
                                                             freq_time_series_data=sel_time_series, selection=None)

            when_results[when_results < 0] = 0  # Bottom Threshold (Remove negative correlation)
            when_elections = np.sum(when_results, axis=1)  # collapse across channels
            # onset_predictions = np.argmax(when_elections[:, expected_before:expected_after], axis=-1)

            syll_all_elections.append(when_elections)
            print('expected_before: ', expected_before)

        super_all_elections.append(syll_all_elections)

    return super_all_elections

def naive_branch_analysis_folded(absolute_relative_starts, motif_ledger, all_best_chunk_events, motif_time_series,
                                 time_buffer, label_instruct):
    """ Algorithm to run a naive onset detection/predicition analysis.

    absolute_relative_starts : list | [labels] -> [instances]
        list of all of the relative start times for each label in the order of labels_used
    motif_ledger : dict | { label_num : [Motif_identity]}
        A ledger of which motif a syllable instance occurred during. It is in order the syllables occurred assuming they
        started with the first syllable
    all_best_chunk_events : list | [Freq]->[labels]->(Instances, 1, Channels, Samples)
        Data clipped into the trials that can be used to make the Pearson Templates
    motif_time_series : list | [Frequency]->(Instances, Channels, Time-Steps, Bin Width)
        list of the times overlapping bins for each frequency to do the mini-onset detection
    time_buffer : int
        #Window around true start time

    label_instruct : list
        list of the syllable labels to use for the Particular bird

    :return:
    [syllables]->[Freqs]->[instances]

    elections: | [syllables]->[Freqs]->[folds]->(instances, samples)
    """

    super_all_elections = []

    for sel_syll in label_instruct:

        syll_all_elections = []
        for sel_freq in range(len(all_best_chunk_events)):
            sel_label = sel_syll - 1

            expected_time = int(np.mean(absolute_relative_starts[sel_label - 1] / 30))
            expected_before = expected_time - time_buffer
            expected_after = expected_time + time_buffer
            print(expected_before, 'to', expected_after)

            # Get the Folds used to make the templates in the Onset Detection analysis
            X = np.arange(len(motif_ledger[sel_syll]))
            y = np.ones((len(motif_ledger[sel_syll]),))
            kf = KFold(n_splits=5, shuffle=True, random_state=0)
            kf.get_n_splits(X)

            print(kf)

            all_elections = []

            # Get the Motifs that do Not have the Syllable (For Branch Analysis)

            # syll_ex = np.arange(len(motif_ledger[sel_syll]))  # Index fo all of the Syllable examples
            branch_ex = [x for x in motif_ledger[2] if x not in motif_ledger[sel_syll]]  # Index of Motifs w/o the Syll

            for train_index, test_index in kf.split(X):
                print("TRAIN:", train_index, "TEST:", test_index)  # These are the only Parts I need for the splitting

                # 1.) Get the Templates
                # 1.1) Sub-select Templates using Template Index
                sel_for_template = all_best_chunk_events[sel_freq][sel_label][train_index]
                sel_template = np.squeeze(np.mean(sel_for_template, axis=0))  # 1.2) Take the Mean (Chan, Samples)

                # 3.) Get the Time-Series [syllable](Instances, channels, samples, bin_width)
                # 3.1) Sub-select Motif Time Series using Branch Index
                sel_time_series = motif_time_series[sel_freq][branch_ex]

                # 4.) Run the Test
                # 4.1) Get the Pearson Coefficients
                when_results = cod.get_freq_pearson_coefficients(freq_template_data=sel_template,
                                                                 freq_time_series_data=sel_time_series, selection=None)

                when_results[when_results < 0] = 0  # Bottom Threshold (Remove negative correlation)
                when_elections = np.sum(when_results, axis=1)  # collapse across channels
                # onset_predictions = np.argmax(when_elections[:, expected_before:expected_after], axis=-1)
                all_elections.append(when_elections)
                print('expected_before: ', expected_before)

            syll_all_elections.append(all_elections)

        super_all_elections.append(syll_all_elections)

    return super_all_elections

def create_shuffled_ledger(motif_ledger):
    """Function to create a Functional Ledger of the location of each Motif after collapsing the shuffled folds

    out:
    shuffled_fold_ledger: dict | {syll: [index_of_instances]
        Note: They go from 0 to num_instances. They are the index of the order of the instances of each syllable
            and not the motif identity.
    """
    shuffled_fold_ledger = {}
    for syll in motif_ledger.keys():

        syllable_holder = []  # Make Holder

        # Get the Folds used to make the templates in the Onset Detection analysis
        X = np.arange(len(motif_ledger[syll]))
        y = np.ones((len(motif_ledger[syll]),))
        kf = KFold(n_splits=5, shuffle=True, random_state=0)
        kf.get_n_splits(X)
        for train_index, test_index in kf.split(X):
            #         print("TRAIN:", train_index, "TEST:", test_index)  # These are the only Parts I need for the splitting
            syllable_holder.append(test_index)  # add folds index order to the the list

        collapsed_index_test = cwan.collaspse_folds(syllable_holder)  # collapse folds
        shuffled_fold_ledger[syll] = collapsed_index_test  # enter entry for current syllable

    return shuffled_fold_ledger

def create_correct_shuffled_ledger(motif_ledger, shuffled_fold_ledger):
    """ function to correct the index value using the shuffled_fold_ledger

    Parameters
    ----------
    motif_ledger : dict | { label_num : [Motif_identity]}
        A ledger of which motif a syllable instance occurred during. It is in order the syllables occurred assuming they
        started with the first syllable
    shuffled_fold_ledger: dict | {syll: [index_of_instances]
        Note: They go from 0 to num_instances. They are the index of the order of the instances of each syllable
            and not the motif identity.

    Returns
    -------
    correct_shuff_ledger :

    """

    # Make a function to correct the index value using the shuffled_fold_ledger

    correct_shuff_ledger = {}
    for syll in motif_ledger.keys():
        holder = []

        for shuff_index in shuffled_fold_ledger[syll]:
            holder.append(motif_ledger[syll][shuff_index])

        correct_shuff_ledger[syll] = np.asarray(holder)
    return correct_shuff_ledger

def organize_onset_prediction(national_election: list, correct_shuff_ledger: dict, full_motif_index: list,
                              full_syllables: list, ):
    """Make an Index of all of the Full Motifs

    Parameters
    ----------
    national_election : list | [syllables]->[folds]->(instances, samples)
         The time-series of all Voting Results for each syllable
    ///
    abs_rel_starts : list | [syllables]->-[folds]>(instances,)
    ///
    correct_shuff_ledger : dict | {syll: [index_of_instances]
        Dictionary of the correct motif index each syllable occurs during in the same order that results from collapsing
        the random shuffled validation folds. This value is the result of using two helper functions.
    full_motif_index : array | (instances,)
        Index of all of the Motifs that have all of the syllables prior to the branch point of interest

    full_syllables : list | [syllables]
        list of syllables prior to the branch point, if there is a branch point, these syllables use the same coding as
        the motif_ledger.

    Returns
    -------
    prediction_timeseries : ndarray | (instances, syllables, samples)
        The prediction/confidence timeseries organized into a more usable format

    abs_rel_starts: ndarray | (instances, syllables,)

    make empty array to hold all of the relevant trials
    """

    # Make empty array to house organized data
    prediction_timeseries = np.zeros((len(full_motif_index), len(full_syllables),
                                      np.shape(national_election[0][0])[-1]))

    # For each index of syllables and prediction results
    for syll_index, (syll, syll_inst_data) in enumerate(zip(full_syllables, national_election)):

        collapsed_elections_data = cwan.collaspse_folds(syll_inst_data)  # Collapse Folds
        viable_motifs_index = correct_shuff_ledger[syll]  # get the Indexes of Motif that have all syllables

        # For each full Motif, get the specified syllable's data
        for motif_index, motif_inst in enumerate(full_motif_index):
            correct_shuffled_index = np.where(viable_motifs_index == motif_inst)  # Find motif index in shuffled ledger
            prediction_timeseries[motif_index, syll_index, :] = collapsed_elections_data[correct_shuffled_index, :]

    return prediction_timeseries

def organize_onset_abs_starts(abs_rel_starts: list, motif_ledger: dict, full_motif_index: list,
                              full_syllables: list, ):
    """Make an Index of all of the Full Motifs

    Parameters
    ----------
    abs_rel_starts : list | [syllables]->[folds]->(instances, samples)
         The time-series of all Voting Results for each syllable
    ///
    abs_rel_starts : list | [syllables]->(instances,)
    ///
    motif_ledger : dict | {syll: [index_of_instances]
       [INSERT HERE]

    full_motif_index : array | (instances,)
        Index of all of the Motifs that have all of the syllables prior to the branch point of interest

    full_syllables : list | [syllables]
        list of syllables prior to the branch point, if there is a branch point, these syllables use the same coding as
        the motif_ledger.

    Returns
    -------
    prediction_timeseries : ndarray | (instances, syllables, samples)
        The prediction/confidence timeseries organized into a more usable format

    abs_rel_starts: ndarray | (instances, syllables,)

    make empty array to hold all of the relevant trials
    """

    # Make empty array to house organized data
    prediction_timeseries = np.zeros((len(full_motif_index), len(full_syllables)))

    # For each index of syllables and prediction results
    for syll_index, (syll, relative_starts) in enumerate(zip(full_syllables, abs_rel_starts)):

        # collapsed_elections_data = cwan.collaspse_folds(relative_starts)  # Collapse Folds
        viable_motifs_index = np.asarray(motif_ledger[syll])  # get the Indexes of Motif that have all syllables

        # For each full Motif, get the specified syllable's data
        for motif_index, motif_inst in enumerate(full_motif_index):
            correct_shuffled_index = np.where(viable_motifs_index == motif_inst)  # Find motif index in shuffled ledger
            prediction_timeseries[motif_index, syll_index] = relative_starts[correct_shuffled_index]

    return prediction_timeseries

def organize_branch_prediction(branch_elections: list, motif_ledger: dict, branch_motif_index: list,
                               branch_syllables: list):
    """Make an Index of all of the Full Motifs

    Parameters
    ----------
    branch_elections : list | [syllables]->[folds]->(instances, samples)
         The time-series of all Voting Results for each syllable
    ///
    abs_rel_starts : list | [syllables]->-[folds]>(instances,)
    ///
    motif_ledger : dict | { label_num : [Motif_identity]}
        A ledger of which motif a syllable instance occurred during. It is in order the syllables occurred assuming they
        started with the first syllable
    branch_motif_index : array | (instances,)
        Index of all of the Motifs that branch at the point of focus

    branch_syllables : list | [syllables]
        list of syllables prior to the branch point, if there is a branch point, these syllables use the same coding as
        the motif_ledger.

    Returns
    -------
    branch_pred_timeseries : ndarray | (instances, syllables, folds, samples)
        The prediction/confidence timeseries for branch syllables organized into a more usable format

    abs_rel_starts: ndarray | (instances, syllables,)

    make empty array to hold all of the relevant trials
    """

    # Make empty array to house organized data
    branch_pred_timeseries = np.zeros((len(branch_motif_index), len(branch_syllables), 5,
                                       np.shape(branch_elections[0][0])[-1]))

    # For each index of syllables and prediction results
    for syll_index, (syll, syll_inst_data) in enumerate(zip(branch_syllables, branch_elections)):

        branch_elections_data = np.asarray(syll_inst_data)  # Shape: [folds]-> (instances, samples)

        # Create a Index of Motifs w/o the Syllable relavent to the branch point of interest
        viable_motifs_index = np.asarray([x for x in motif_ledger[2] if x not in motif_ledger[syll]])

        # For each full Motif, get the specified syllable's data
        for motif_index, motif_inst in enumerate(branch_motif_index):
            correct_shuffled_index = np.where(viable_motifs_index == motif_inst)  # Find motif index in shuffled ledger
            branch_pred_timeseries[motif_index, syll_index, :, :] = np.squeeze(
                np.squeeze(branch_elections_data[:, correct_shuffled_index, :]))

    return branch_pred_timeseries

def organize_branch_prediction_single(branch_elections: list, motif_ledger: dict, branch_motif_index: list,
                                      branch_syllables: list):
    """Make an Index of all of the Full Motifs

    Parameters
    ----------
    branch_elections : list | [folds]->(instances, samples)
         The time-series of all Voting Results for each syllable
    ///
    abs_rel_starts : list | [folds]>(instances,)
    ///
    motif_ledger : dict | { label_num : [Motif_identity]}
        A ledger of which motif a syllable instance occurred during. It is in order the syllables occurred assuming they
        started with the first syllable
    branch_motif_index : array | (instances,)
        Index of all of the Motifs that branch at the point of focus

    branch_syllables : list | [syllables]
        list of syllables prior to the branch point, if there is a branch point, these syllables use the same coding as
        the motif_ledger.

    Returns
    -------
    branch_pred_timeseries : ndarray | (instances, syllables, folds, samples)
        The prediction/confidence timeseries for branch syllables organized into a more usable format

    abs_rel_starts: ndarray | (instances, syllables,)

    make empty array to hold all of the relevant trials
    """

    # Make empty array to house organized data
    branch_pred_timeseries = np.zeros((len(branch_motif_index), len(branch_syllables), 5,
                                       np.shape(branch_elections[0][0])[-1]))

    # For each index of syllables and prediction results
    #     for syll_index, (syll, syll_inst_data) in enumerate(zip(branch_syllables, branch_elections)):

    branch_elections_data = np.asarray(branch_elections)  # Shape: [folds]-> (instances, samples)

    # Create a Index of Motifs w/o the Syllable relavent to the branch point of interest
    viable_motifs_index = np.asarray([x for x in motif_ledger[2] if x not in motif_ledger[branch_syllables[0]]])

    # For each full Motif, get the specified syllable's data
    for motif_index, motif_inst in enumerate(branch_motif_index):
        correct_shuffled_index = np.where(viable_motifs_index == motif_inst)  # Find motif index in shuffled ledger
        branch_pred_timeseries[motif_index, 0, :, :] = np.squeeze(
            np.squeeze(branch_elections_data[:, correct_shuffled_index, :]))

    return branch_pred_timeseries

def make_nested_list_of_results(absolute_relative_starts, national_election, branch_election_collapsed, syll_focus,
                                time_buff=50):
    syll_with = []
    syll_without = []

    for sel_fold in range(5):
        #     print(sel_fold)
        expected_time = int(np.mean(absolute_relative_starts[syll_focus] / 30))
        expected_before = expected_time - time_buff
        expected_after = expected_time + time_buff

        # for Syllable 6
        result_with = np.max(national_election[syll_focus][sel_fold][:, expected_before:expected_after], axis=-1)
        result_without = np.max(branch_election_collapsed[sel_fold][:, expected_before:expected_after],
                                axis=-1)

        syll_with.append(result_with)
        syll_without.append(result_without)

    syll_with = np.asarray(syll_with)
    syll_without = np.asarray(syll_without)

    return syll_with, syll_without

def make_array_of_max_results(absolute_relative_starts, national_election, branch_election_collapsed, syll_focus,
                              time_buff=50):
    syll_with = []
    syll_without = []

    for sel_fold in range(5):
        #     print(sel_fold)
        expected_time = int(np.mean(absolute_relative_starts[syll_focus] / 30))
        expected_before = expected_time - time_buff
        expected_after = expected_time + time_buff

        # for Syllable 6
        result_with = np.max(national_election[syll_focus][sel_fold][:, expected_before:expected_after], axis=-1)
        result_without = np.max(branch_election_collapsed[sel_fold][:, expected_before:expected_after],
                                axis=-1)

        syll_with.extend(result_with)
        syll_without.extend(result_without)

    syll_with = np.asarray(syll_with)
    syll_without = np.asarray(syll_without)

    return syll_with, syll_without

def get_branch_statistics(syll_with_list, syll_without_list):
    all_p_values = []
    for with_syllable, without_syllable in zip(syll_with_list, syll_without_list):
        t_stat, p_value = scipy.stats.ttest_ind(with_syllable, without_syllable, equal_var=False, alternative="greater")
        print(p_value)
        all_p_values.append(p_value)
    return all_p_values

def run_branch_point_analysis_script(bird_id, session):
    # Import the Data
    zdata = ImportData(bird_id=bird_id, session=session)

    # Get the Bird Specific Machine Learning Meta Data
    bad_channels = all_bad_channels[bird_id]
    drop_temps = all_drop_temps[bird_id]
    when_label_instructions = cwan.all_when_label_instructions[bird_id]
    time_buffer = 50

    if session == 'day-2016-09-09':
        bin_widths = [185, 155, 120, 50, 120]
        offsets = [-40, -25, 0, 0, 0]
    else:
        # Get the Best Parameters for Bindwidth and Offset
        bin_widths = best_bin_width[session]
        offsets = best_offset[session]

    # Reshape Handlabels into Useful Format
    chunk_labels_list, chunk_onsets_list = fet.get_chunk_handlabels(handlabels_list=zdata.song_handlabels)

    # Set the Frequency Bands to Be Used for Feature Extraction
    fc_lo = [4, 8, 25, 30, 50]
    fc_hi = [8, 12, 35, 50, 70]

    # Pre-Process the Data # TODO: Make this a function
    # [Freqs]->[Chunks]->(1, channels, samples)

    freq_pred_data = []
    for low, high in zip(fc_lo, fc_hi):
        pred_data = hbp.feature_extraction_chunk(neural_chunks=zdata.song_neural,
                                                 fs=1000,
                                                 l_freqs=[low],
                                                 h_freqs=[high],
                                                 hilbert=None,
                                                 bad_channels=bad_channels,
                                                 norm=False,
                                                 drop_bad=True,
                                                 verbose=True)
        freq_pred_data.append(pred_data)

    # Get the Bird Specific label Instructions
    label_instructions = all_label_instructions[bird_id]  # get this birds default label instructions

    times_of_interest = fet.label_extractor(all_labels=chunk_labels_list,
                                            starts=chunk_onsets_list[0],
                                            label_instructions=label_instructions)

    # [Freq]->[labels]->(Instances, 1, Channels, Samples)
    all_best_chunk_events_test = cod.clip_instances_for_templates(freq_preprocessed_data=freq_pred_data,
                                                                  all_offsets=offsets,
                                                                  all_bin_widths=bin_widths,
                                                                  label_events=times_of_interest)

    motif_length = cwan.motif_length_dict[bird_id]

    # [Frequency]->(Instances, Channels, Time-Steps, Bin Width)
    motif_time_series = cod.clip_motif_time_series(freq_preprocessed_data=freq_pred_data,
                                                   all_offsets=offsets,
                                                   all_bin_widths=bin_widths,
                                                   motif_start_times=times_of_interest[0],
                                                   motif_length=motif_length)

    # Create instance of the Context Labels Class
    testclass = birds_context_obj(bird_id=bird_id)

    # Get the Context Array for the Day's Data
    test_context = testclass.get_all_context_index_arrays(chunk_labels_list)

    # Make Index of the Motif the syllable belongs to
    # list of the labels (syllable 2 : 2)
    # when_label_instructions = cod.fix_label_instructions(label_instructions)

    #  {label: [motif indexes]}
    motif_ledger = cod.get_motif_identifier(focus=when_label_instructions,
                                            context=test_context,
                                            labels=chunk_labels_list)

    # Get all of the Relative Starts of the Syllables in their Respective Motifs
    # [label]-> [Instances]
    organized_times_of_interest = cod.organize_absolute_times(times_of_interest)

    # Get the Relative Start Times (in 30KHz Samples) [Need to divide by 30 to get the onset in ms]
    absolute_relative_starts = cod.get_all_relative_starts(first_starts=organized_times_of_interest[0],
                                                           other_starts=organized_times_of_interest[1:],
                                                           all_other_index=motif_ledger,
                                                           labels_used=when_label_instructions)

    cwan._remove_unusable_syllable(all_best_chunk_events=all_best_chunk_events_test, motif_ledger=motif_ledger,
                                   label_instruct=when_label_instructions)  # Remove Bad Syllable Instances

    cwan._fix_the_ledger(motif_ledger)  # New Function to remove 'bad' from the ledger

    label_instruct = all_label_instruct[bird_id]

    branch_elections_test = naive_branch_analysis_folded(absolute_relative_starts=absolute_relative_starts,
                                                         motif_ledger=motif_ledger,
                                                         all_best_chunk_events=all_best_chunk_events_test,
                                                         motif_time_series=motif_time_series,
                                                         time_buffer=time_buffer,
                                                         label_instruct=label_instruct)

    # Collapse Across Frequencies
    branch_election_collapsed = cwan.get_national_elections(elections=branch_elections_test)

    # Make a maping from the audio to the Corresponding Prediction Trials
    shuffled_fold_ledger = create_shuffled_ledger(motif_ledger=motif_ledger)

    # Make a ledger/index to correct the index value using the shuffled_fold_ledger
    correct_shuff_ledger = create_correct_shuffled_ledger(motif_ledger=motif_ledger,
                                                          shuffled_fold_ledger=shuffled_fold_ledger)
    ######
    # Save the National Elections
    national_election = _load_pckl_data(data_name="national_election_run_2", bird_id=bird_id,
                                        session=session,
                                        source=onset_detection_path, verbose=True)

    # Save the Motif Ledger
    absolute_relative_starts = _load_pckl_data(data_name="absolute_relative_starts_run_2", bird_id=bird_id,
                                               session=session, source=onset_detection_path, verbose=True)
    ######

    # Organize the Correlation trace for all of the full syllables
    full_motif_organized = organize_onset_prediction(national_election=national_election,
                                                     correct_shuff_ledger=correct_shuff_ledger,
                                                     full_motif_index=motif_ledger[label_instruct[-1]],
                                                     full_syllables=motif_ledger.keys())

    # Organize the Start Times for all of the full syllables
    full_motif_starts = organize_onset_abs_starts(abs_rel_starts=absolute_relative_starts,
                                                  motif_ledger=motif_ledger,
                                                  full_motif_index=motif_ledger[label_instruct[-1]],
                                                  full_syllables=motif_ledger.keys())

    # Write a Branch Point for The Script to handle the two main birds with Branch Points
    # This needs to save this data as a pickle so that I can import it and use it to both run Statistics and Make Figures

    # Explicitly set three if statements for each bird:
    if bird_id == "z017":
        # Branch 1: Full Syllable

        branch_1 = [x for x in motif_ledger[5] if x not in motif_ledger[6]]  # Index of Motifs w/o the Syll 7

        # Correlation trace for all of the Syllables prior to the branch point
        branch_organized_1 = organize_onset_prediction(national_election=national_election,
                                                       correct_shuff_ledger=correct_shuff_ledger,
                                                       full_motif_index=branch_1, full_syllables=[2, 3, 4, 5])

        # Correlation trace for all of the branch syllable(s)
        branch_predictions_1 = organize_branch_prediction(branch_elections=branch_election_collapsed,
                                                          motif_ledger=motif_ledger, branch_motif_index=branch_1,
                                                          branch_syllables=[6, 7])

        branch_2 = np.asarray(
            [x for x in motif_ledger[6] if x not in motif_ledger[7]])  # Index of Motifs w/o the Syll 7

        # Correlation trace for all of the Syllables prior to the branch point
        branch_organized_2 = organize_onset_prediction(national_election=national_election,
                                                       correct_shuff_ledger=correct_shuff_ledger,
                                                       full_motif_index=branch_2,
                                                       full_syllables=[2, 3, 4, 5, 6])

        # Correlation trace for all of the branch syllable(s)
        branch_predictions_2 = organize_branch_prediction_single(branch_elections=branch_election_collapsed[1],
                                                                 motif_ledger=motif_ledger,
                                                                 branch_motif_index=branch_2,
                                                                 branch_syllables=[7])

        # Get Branch 1: absolute starts for syllables before the branch point
        branch_starts_1 = organize_onset_abs_starts(abs_rel_starts=absolute_relative_starts,
                                                    motif_ledger=motif_ledger,
                                                    full_motif_index=branch_1,
                                                    full_syllables=[2, 3, 4, 5])

        # Get Branch 2: absolute starts for syllables before the branch point
        branch_starts_2 = organize_onset_abs_starts(abs_rel_starts=absolute_relative_starts,
                                                    motif_ledger=motif_ledger,
                                                    full_motif_index=branch_2,
                                                    full_syllables=[2, 3, 4, 5, 6])

        ## Run Stats
        # Make a Nested List to run Statistics for Sylable 6
        with_nest6, without_nest6 = make_nested_list_of_results(absolute_relative_starts=absolute_relative_starts,
                                                                national_election=national_election,
                                                                branch_election_collapsed=branch_election_collapsed[0],
                                                                syll_focus=4, time_buff=time_buffer)

        syll_6_stats = get_branch_statistics(syll_with_list=with_nest6, syll_without_list=without_nest6)

        max_with_6, max_without_6 = make_array_of_max_results(absolute_relative_starts=absolute_relative_starts,
                                                              national_election=national_election,
                                                              branch_election_collapsed=branch_election_collapsed[0],
                                                              syll_focus=4, time_buff=time_buffer)

        # Make a Nested List to run Statistics for Sylable 7
        with_nest7, without_nest7 = make_nested_list_of_results(absolute_relative_starts=absolute_relative_starts,
                                                                national_election=national_election,
                                                                branch_election_collapsed=branch_election_collapsed[1],
                                                                syll_focus=5, time_buff=time_buffer)

        syll_7_stats = get_branch_statistics(syll_with_list=with_nest7, syll_without_list=without_nest7)

        max_with_7, max_without_7 = make_array_of_max_results(absolute_relative_starts=absolute_relative_starts,
                                                              national_election=national_election,
                                                              branch_election_collapsed=branch_election_collapsed[1],
                                                              syll_focus=5, time_buff=time_buffer)

        ## Save the Componenets for the Summary Figure
        # Save the DataFrame of Onset Results
        _save_pckl_data(data=branch_organized_1, data_name="branch_organized_1", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)
        _save_pckl_data(data=branch_predictions_1, data_name="branch_predictions_1", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)
        _save_pckl_data(data=branch_organized_2, data_name="branch_organized_2", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)
        _save_pckl_data(data=branch_predictions_2, data_name="branch_predictions_2", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)

        _save_pckl_data(data=branch_starts_1, data_name="branch_starts_1", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)
        _save_pckl_data(data=branch_starts_2, data_name="branch_starts_2", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)

        # Stats & Analysis
        _save_pckl_data(data=syll_6_stats, data_name="syll_6_stats", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)
        _save_pckl_data(data=max_with_6, data_name="max_with_6", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)
        _save_pckl_data(data=max_without_6, data_name="max_without_6", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)

        _save_pckl_data(data=syll_7_stats, data_name="syll_7_stats", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)
        _save_pckl_data(data=max_with_7, data_name="max_with_7", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)
        _save_pckl_data(data=max_without_7, data_name="max_without_7", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)

    elif bird_id == "z020":
        # Branch 1: Full Syllable

        branch_1 = [x for x in motif_ledger[3] if x not in motif_ledger[4]]  # Index of Motifs w/o the Syll 7

        # Correlation trace for all of the Syllables prior to the branch point
        branch_organized_1 = organize_onset_prediction(national_election=national_election,
                                                       correct_shuff_ledger=correct_shuff_ledger,
                                                       full_motif_index=branch_1, full_syllables=[2, 3])

        # Correlation trace for all of the branch syllable(s)
        branch_predictions_1 = organize_branch_prediction(branch_elections=branch_election_collapsed,
                                                          motif_ledger=motif_ledger, branch_motif_index=branch_1,
                                                          branch_syllables=[4])

        # Get Branch 1: absolute starts for syllables before the branch point
        branch_starts_1 = organize_onset_abs_starts(abs_rel_starts=absolute_relative_starts,
                                                    motif_ledger=motif_ledger,
                                                    full_motif_index=branch_1,
                                                    full_syllables=[2, 3])

        ## Run Statistics
        # Make a Nested List to run Statistics for Sylable 4 (syll 4 => syll_focus = 2)
        with_nest4, without_nest4 = make_nested_list_of_results(absolute_relative_starts=absolute_relative_starts,
                                                                national_election=national_election,
                                                                branch_election_collapsed=branch_election_collapsed[0],
                                                                syll_focus=2, time_buff=time_buffer)

        syll_4_stats = get_branch_statistics(syll_with_list=with_nest4, syll_without_list=without_nest4)

        max_with_4, max_without_4 = make_array_of_max_results(absolute_relative_starts=absolute_relative_starts,
                                                              national_election=national_election,
                                                              branch_election_collapsed=branch_election_collapsed[0],
                                                              syll_focus=2, time_buff=time_buffer)

        _save_pckl_data(data=branch_organized_1, data_name="branch_organized_1", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)
        _save_pckl_data(data=branch_predictions_1, data_name="branch_predictions_1", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)
        _save_pckl_data(data=branch_starts_1, data_name="branch_starts_1", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)

        # Stats & Analysis
        _save_pckl_data(data=syll_4_stats, data_name="syll_4_stats", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)
        _save_pckl_data(data=max_with_4, data_name="max_with_4", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)
        _save_pckl_data(data=max_without_4, data_name="max_without_4", bird_id=bird_id,
                        session=session,
                        destination=branch_analysis_path, make_parents=True, verbose=True)

    elif bird_id == "z007":
        pass

    else:
        print("this subject was not accounted for in this script")

    # Save the organized results for the Full Motifs
    _save_pckl_data(data=full_motif_organized, data_name="full_motif_organized", bird_id=bird_id,
                    session=session,
                    destination=branch_analysis_path, make_parents=True, verbose=True)
    _save_pckl_data(data=full_motif_starts, data_name="full_motif_starts", bird_id=bird_id,
                    session=session,
                    destination=branch_analysis_path, make_parents=True, verbose=True)


from torch import nn
import torch
import torch.nn.functional as F
from torch.autograd import Variable
import torchvision
from torchvision.ops import nms
from nts_net import resnet
import numpy as np
from nts_net.anchors import generate_default_anchor_maps
from nts_net.config import CAT_NUM, PROPOSAL_NUM, test_model

class ProposalNet(nn.Module):
    def __init__(self):
        super(ProposalNet, self).__init__()
        self.down1 = nn.Conv2d(2048, 128, 3, 1, 1)
        self.down2 = nn.Conv2d(128, 128, 3, 2, 1)
        self.down3 = nn.Conv2d(128, 128, 3, 2, 1)
        self.ReLU = nn.ReLU()
        self.tidy1 = nn.Conv2d(128, 6, 1, 1, 0)
        self.tidy2 = nn.Conv2d(128, 6, 1, 1, 0)
        self.tidy3 = nn.Conv2d(128, 9, 1, 1, 0)

    def forward(self, x):
        batch_size = x.size(0)
        d1 = self.ReLU(self.down1(x))
        d2 = self.ReLU(self.down2(d1))
        d3 = self.ReLU(self.down3(d2))
        t1 = self.tidy1(d1).view(batch_size, -1)
        t2 = self.tidy2(d2).view(batch_size, -1)
        t3 = self.tidy3(d3).view(batch_size, -1)
        return torch.cat((t1, t2, t3), dim=1)

class attention_net(nn.Module):
    def __init__(self, topN=6, device='cpu', num_classes=200):
        super(attention_net, self).__init__()
        self.device = device
        self.pretrained_model = resnet.resnet50(pretrained=True)
        self.pretrained_model.avgpool = nn.AdaptiveAvgPool2d(1)
        self.pretrained_model.fc = nn.Linear(512 * 4, num_classes)
        self.proposal_net = ProposalNet()
        self.topN = topN
        self.concat_net = nn.Linear(2048 * (CAT_NUM + 1), num_classes)
        self.partcls_net = nn.Linear(512 * 4, num_classes)
        _, edge_anchors, _ = generate_default_anchor_maps()
        self.pad_side = 224
        self.edge_anchors = (edge_anchors + 224).astype(np.int)

    def forward(self, x):
        resnet_out, rpn_feature, feature = self.pretrained_model(x)
        x_pad = F.pad(x, (self.pad_side, self.pad_side, self.pad_side, self.pad_side), mode='constant', value=0)
        batch = x.size(0)
        # we will reshape rpn to shape: batch * nb_anchor
        rpn_score = self.proposal_net(rpn_feature.detach())
        top_n_index_list = []
        top_n_prob_list = []
        top_n_coordinates = []
        edge_anchors_copy = torch.tensor(self.edge_anchors, dtype=torch.float32).to(self.device)
        zero_tensor = torch.zeros(rpn_score.size()[1], 1)

        for i in range(batch):
            rpn_score_reshape = rpn_score[i]
            nms_output_idx = nms(edge_anchors_copy, rpn_score_reshape, iou_threshold=0.25)
            nms_output_score = torch.gather(rpn_score_reshape, 0, nms_output_idx)
            nms_output_anchors = torch.index_select(edge_anchors_copy, 0, nms_output_idx)
            top_n_result = torch.topk(nms_output_score, self.topN)
            top_n_anchors = torch.index_select(nms_output_anchors, 0, top_n_result[1])
            y0_1, x0_1, y1_1, x1_1 = torch.split(top_n_anchors, 1, dim=1)
            top_n_anchors_1 = torch.cat([x0_1, y0_1, x1_1, y1_1], dim=1)
            top_n_index_origin = torch.index_select(nms_output_idx, 0, top_n_result[1])
            top_n_index_list.append(top_n_index_origin)
            top_n_prob_list.append(top_n_result[0])
            top_n_coordinates.append(top_n_anchors_1)

        top_n_index = torch.stack(top_n_index_list)
        top_n_prob = torch.stack(top_n_prob_list)
        part_imgs = torch.zeros([batch, self.topN, 3, 224, 224])
        for i in range(batch):
            for j in range(self.topN):
                y0 = top_n_coordinates[i][j][1].long()
                x0 = top_n_coordinates[i][j][0].long()
                y1 = top_n_coordinates[i][j][3].long()
                x1 = top_n_coordinates[i][j][2].long()
                part_imgs[i:i + 1, j] = F.interpolate(x_pad[i:i + 1, :, y0:y1, x0:x1], size=(224, 224), mode='bilinear',
                                                      align_corners=True)
        part_imgs = part_imgs.view(batch * self.topN, 3, 224, 224).to(self.device)
        _, _, part_features = self.pretrained_model(part_imgs.detach())
        part_feature = part_features.view(batch, self.topN, -1)
        part_feature = part_feature[:, :CAT_NUM, ...].contiguous()
        part_feature = part_feature.view(batch, -1)
        # concat_logits have the shape: B*200
        concat_out = torch.cat([part_feature, feature], dim=1)
        concat_logits = self.concat_net(concat_out)
        raw_logits = resnet_out
        # part_logits have the shape: B*N*200
        part_logits = self.partcls_net(part_features).view(batch, self.topN, -1)
        return top_n_coordinates, concat_out, raw_logits, concat_logits, part_logits, top_n_index, top_n_prob

def list_loss(logits, targets):
    temp = F.log_softmax(logits, -1)
    loss = [-temp[i][targets[i].item()] for i in range(logits.size(0))]
    return torch.stack(loss)

def ranking_loss(score, targets, proposal_num=PROPOSAL_NUM):
    loss = Variable(torch.zeros(1).cuda())
    batch_size = score.size(0)
    for i in range(proposal_num):
        targets_p = (targets > targets[:, i].unsqueeze(1)).type(torch.cuda.FloatTensor)
        pivot = score[:, i].unsqueeze(1)
        loss_p = (1 - pivot + score) * targets_p
        loss_p = torch.sum(F.relu(loss_p))
        loss += loss_p
    return loss / batch_size

def ntsnet(model_path, device='cuda'):
    net = attention_net(topN=PROPOSAL_NUM, device=device)
    ckpt = torch.load(model_path)
    net.load_state_dict(ckpt['net_state_dict'])
    print('test_acc:', ckpt['test_acc'])
    return net

if __name__ == '__main__':
    net = ntsnet(test_model)


#coding=utf-8

############################################################################
#
# Copyright (c) 2018 ICT MCG Group, Inc. All Rights Reserved
#
###########################################################################
"""
Brief:
任务：单微博的谣言检测
方法：实现文本模态和视觉模态的融合
1. baseline:拼接
2. 每个post动态训练权重

Authors: qipeng(@ict.ac.cn)
Date:    2018/09/19 18:41:48
File:    fusion.py
"""
import sys

from keras.models import Sequential, Model, load_model
from keras.layers import Dense, Activation, Embedding, Dropout, LSTM, Merge, Lambda, Input, merge, Flatten, Bidirectional, TimeDistributed, multiply
from keras.layers.core import Reshape, RepeatVector, Permute
from keras.optimizers import SGD, Adam
from keras.regularizers import l2
from attention_layer import AttentionLayer3
import tensorflow as tf
from keras import backend as K
import keras
import numpy as np
import pandas as pd
from sklearn import metrics
from sklearn.model_selection import KFold,GroupKFold,StratifiedKFold

import time
import os
os.environ['CUDA_VISIBLE_DEVICES']='0,1,2,3'

import pickle

## PATH
model_path = './data/model_weights_'+ time.strftime("%y-%m-%d-%H-%M-%S", time.localtime()) +'.h5'
history_path = './data/train_logs_' + time.strftime("%y-%m-%d-%H-%M-%S", time.localtime()) + '.txt'
weights_path = './data/weights_' + time.strftime("%y-%m-%d-%H-%M-%S", time.localtime()) +'.txt'
score_path = './data/score_' + time.strftime("%y-%m-%d-%H-%M-%S", time.localtime()) + '.txt'

## WEIGHTS
maxlen = 80
batch_size = 256
epoch_size = 1
embedding_dim = 32
hidden_num = 16

## FEATURES，注意两种特征要在一个scale上
# text feature
test_neg = np.load('./numpy_data/test_neg.npy')
test_pos = np.load('./numpy_data/test_pos.npy')
train_neg = np.load('./numpy_data/train_neg.npy')
train_pos = np.load('./numpy_data/train_pos.npy')

# visual feature
auxi_test_neg_vis = np.load('./features/vis_simplemodel/test_rumor_vis.npy')
auxi_test_pos_vis = np.load('./features/vis_simplemodel/test_nonrumor_vis.npy')
auxi_train_neg_vis = np.load('./features/vis_simplemodel/train_rumor_vis.npy')
auxi_train_pos_vis = np.load('./features/vis_simplemodel/train_nonrumor_vis.npy')

X_train = np.concatenate((train_pos, train_neg))
X_train_auxi_vis = np.concatenate((auxi_train_pos_vis, auxi_train_neg_vis))
train_label = np.concatenate((np.zeros(train_pos.shape[0]),np.ones(train_neg.shape[0])))
y_train = pd.get_dummies(train_label).as_matrix()

X_test = np.concatenate((test_pos, test_neg))
X_test_auxi_vis = np.concatenate((auxi_test_pos_vis, auxi_test_neg_vis))
test_label = np.concatenate((np.zeros(test_pos.shape[0]),np.ones(test_neg.shape[0])))
y_test = pd.get_dummies(test_label).as_matrix() # 单位矩阵

print("X_train Shape:", X_train.shape)
print("y_train Shape:", y_train.shape)
print("X_test Shape:", X_test.shape)
print("y_test Shape:", y_test.shape)

# output得到的是两个类别的概率，需要比较大小确定最终类别
# 把谣言类当做主类
def convert_label(label):
    res = []
    label = label.tolist()
    for ele in label:
        if ele[0] > ele[1]:
            res.append(0)
        else:
            res.append(1)
    return res

#把非谣言当做主类
def convert_label_nonrumor(label):
    res = []
    label = label.tolist()
    for ele in label:
        if ele[0] > ele[1]:
            res.append(1)
        else:
            res.append(0)
    return res

def build_base_model():

    ## text modality
    input_shape = (maxlen, embedding_dim)
    input_data = Input(input_shape, name = 'main_input')

    bilstm_out = Bidirectional(LSTM(hidden_num, return_sequences = True), merge_mode = 'concat')(input_data)
    bilstm_out = Bidirectional(LSTM(hidden_num, return_sequences = True), merge_mode = 'concat')(bilstm_out)

    # attention
    att_out = AttentionLayer3(name='attentionLayer')(bilstm_out)
    # [None, 20, 1]
    att_rep = Lambda(lambda x : K.repeat_elements(x, hidden_num * 2, axis=2))(att_out)
    merge_out = multiply([att_rep, bilstm_out])
    sum_out = Lambda(lambda x: K.sum(x, axis=1))(merge_out)

    ## visual modality
    auxi_input_vis = Input(shape = (auxi_test_pos_vis.shape[1],),name = 'auxi_input_vis')
    x_vis = Dense(hidden_num*2, activation = 'relu', kernel_regularizer = l2(0.005))(auxi_input_vis)

    x = keras.layers.concatenate([sum_out, x_vis]) # fusion way

    # x = Dense(64, activation = 'relu')(x)
    x = Dense(32, activation = 'relu', kernel_regularizer = l2(0.005))(x)

    main_out = Dense(2, activation = 'softmax', name = 'main_output')(x)

    net = Model(input = [input_data, auxi_input_vis], output =  main_out)
    return net

def build_fusion_model():

    ########### text modality ##########
    # shape(None, maxlen, embedding_dim)
    input_text = Input(shape = (maxlen, embedding_dim), name = 'input_text')
    # (None, max_len, hidden_num * 2)
    # Q: 输出是None, None, hidden_num*2，可能是因为具体长度无法确定，数据预处理的时候做补齐和截断了吗
    bilstm_out1 = Bidirectional(LSTM(hidden_num, return_sequences = True), merge_mode = 'concat')(input_text)
    bilstm_out2 = Bidirectional(LSTM(hidden_num, return_sequences = True), merge_mode = 'concat')(bilstm_out1)
    # attention (None, max_len, 1)
    att_out = AttentionLayer3(name='attentionLayer')(bilstm_out2)
    # (None, max_len, hidden_num * 2)
    att_rep = Lambda(lambda x : K.repeat_elements(x, hidden_num * 2, axis=2))(att_out)
    # (None, max_len, hidden_num * 2)
    merge_out = multiply([att_rep, bilstm_out2])
    # (None, hidden_num * 2)
    x_text = Lambda(lambda x: K.sum(x, axis=1))(merge_out)

    ########### visiual modality ##########
    input_vis = Input(shape = (auxi_test_pos_vis.shape[1],),name = 'input_vis')
    # (None, hidden_num * 2)
    # x_vis = Dense(hidden_num*2, activation = 'relu', kernal_regularizer = l2(0.03))(input_vis)
    x_vis = Dense(hidden_num*2, activation = 'relu', kernel_regularizer = l2(0.01))(input_vis)

    ########### fusion ##########
    # (None, hidden_num * 4)
    x = keras.layers.concatenate([x_text, x_vis])

    # (None, 2)
    weights = Dense(2, activation='softmax', name='weights', kernel_regularizer = l2(0.01))(x)
    # (None, 2, 1)
    my_expand_dims = Lambda(lambda xx: K.expand_dims(xx, axis=2))
    weights1 = my_expand_dims(weights)
    # (None, 2, hidden_num *2)
    # weights = K.repeat_elements(weights, hidden_num*2, axis=1) 在函数式模型中调用后端函数需要用lambda进行封装
    my_weights = Lambda(lambda xx: K.repeat_elements(xx, hidden_num*2, axis=2))
    weights2 = my_weights(weights1)
    # print weights2.shape
    # (None, hidden_num *4, 1)
    weights3 = keras.layers.core.Reshape((hidden_num * 4, 1),name='weights_e')(weights2)
    # print weights3.shape
    # (None, hidden_num *4)
    my_sum = Lambda(lambda x: K.sum(x, axis=-1))
    weights4 = my_sum(weights3)
    # print weights4.shape

    # (None, hidden_num*4)
    x_fusion = keras.layers.multiply([x, weights4])

    # (None, 64)
    #x_fusion1 = Dense(64, activation = 'relu')(x_fusion)
    # (None, 32)
    x_fusion2 = Dense(32, activation = 'relu', kernel_regularizer = l2(0.01))(x_fusion)

    # my_dropout = Dropout(0.5)
    # x_fusion2 = my_dropout(x_fusion2)

    # main_out = Dense(2, activation = 'softmax', kernel_regularizer = l2(0.03), name = 'main_output')(x_fusion2)
    main_out = Dense(2, activation = 'softmax', name = 'main_output', kernel_regularizer = l2(0.01))(x_fusion2)
    #print main_out.shape

    model = Model(input = [input_text, input_vis], output =  main_out)
    return model

def compile_model(model):
    print "Training model..."

    print(model.summary())
    adam = Adam(lr = 0.001, beta_1 = 0.9, beta_2 = 0.999, epsilon = 1e-08, decay = 0.0)
    model.compile(loss='categorical_crossentropy', optimizer= adam, metrics=['accuracy'])
    # history = model.fit([X_train, X_train_auxi_vis],y_train, nb_epoch = epoch_size, batch_size = batch_size, validation_data = ([X_test,  X_test_auxi_vis], y_test))
    history = model.fit([X_train, X_train_auxi_vis],y_train, nb_epoch = epoch_size, batch_size =
            batch_size, shuffle = True, validation_split = 0.3)

    hist_logs = history.history
    print "max_val_acc: %f" % max(hist_logs['val_acc'])
    print "Saving history logs..."
    with open(history_path, 'wb') as f:
        pickle.dump(hist_logs, f)

    print "Saving model..."
    model.save(model_path)

def get_weights(model):
    weights_model = Model(inputs=model.input, outputs=model.get_layer('weights').output)
    weights_output = weights_model.predict([X_test, X_test_auxi_vis])
    print "Saving weights..."
    np.savetxt(weights_path, weights_output)

def predict_model(model):
    pre_label = model.predict([X_test, X_test_auxi_vis])

    print('#################')
    print('## NonRumor')
    print('#################')
    pre_label_c = convert_label_nonrumor(pre_label)
    test_label_c = convert_label_nonrumor(y_test)
    acc = metrics.accuracy_score(test_label_c, pre_label_c)
    precision = metrics.precision_score(test_label_c, pre_label_c)
    recall = metrics.recall_score(test_label_c, pre_label_c)
    f1 = metrics.f1_score(test_label_c, pre_label_c)
    print 'Accuracy: ' + str(acc)
    print 'Precision: ' + str(precision)
    print 'Recall: ' + str(recall)
    print 'F1: ' + str(f1)

    print('#################')
    print('## Average')
    print('#################')
    acc = metrics.accuracy_score(test_label_c, pre_label_c)
    precision = metrics.precision_score(test_label_c, pre_label_c, average = 'macro')
    recall = metrics.recall_score(test_label_c, pre_label_c, average = 'macro')
    f1 = metrics.f1_score(test_label_c, pre_label_c, average = 'macro')
    print('Accuracy: ' + str(acc))
    print('Precision: ' + str(precision))
    print('Recall: ' + str(recall))
    print('F1: ' + str(f1))

    print('#################')
    print('## Rumor')
    print('#################')
    pre_label_c = convert_label(pre_label)
    test_label_c = convert_label(y_test)
    acc = metrics.accuracy_score(test_label_c, pre_label_c)
    precision = metrics.precision_score(test_label_c, pre_label_c)
    recall = metrics.recall_score(test_label_c, pre_label_c)
    f1 = metrics.f1_score(test_label_c, pre_label_c)
    print('Accuracy: ' + str(acc))
    print('Precision: ' + str(precision))
    print('Recall: ' + str(recall))
    print('F1: ' + str(f1))

    fpr, tpr, thresholds = metrics.roc_curve(test_label_c, pre_label_c)
    auc = metrics.auc(fpr, tpr)
    print ('AUC: ' + str(auc))

def main():
    if (os.path.exists(model_path)):
        print "Model exists. Loading model..."
        model = load_model(model_path)
    else:
        print "Model doesn't exist. Building model..."
        model = build_fusion_model()
        # model = build_base_model()

    start = time.time()
    compile_model(model)
    get_weights(model)
    predict_model(model)
    end = time.time()
    print "time consumed: %.2f s" % (end - start)

if __name__ == '__main__':
    main()
# vim: set expandtab ts=4 sw=4 sts=4 tw=100:


#!/usr/bin/env python
import re
import hashlib
import random
import subprocess as sp
from common.ascii_art import AsciiImage, AsciiText
import time
import sys
import string
import argparse
import pandas as pd
import numpy as np
import os

def parsey():
    parser = argparse.ArgumentParser(description='Sorting Hat sorts, coders code.')
    parser.add_argument('-train', dest='train', metavar='N', type=str, nargs=1,
                        help='File used to train the sorting hat each year.')
    parser.add_argument('-debug', dest='debug', action='store_true',
                        help='Set debug state true')
    args = parser.parse_args()
    return args

class sortingHat():
    '''This is the Sorting Hat.
    The Sorting Hat uses a "trained" pseudo-random number generator
    seeded with the users name to determine the house placement for the user.
    The Hat needs to be trained with a new set of names each time it is run.'''
    def __init__(self):
        self.debug = False
        self.training = False
        self.set_Constants()
        self.get_flags()  # sets training state and debug state
        if not self.training:
            self.isPM = False
            self.quoteSpecified = False
            self.get_name()
            self.check_name()
            self.sort()
            self.show_result()
        elif self.training:
            self.train()

    def get_flags(self):
        flags = parsey()
        if flags.train is not None:
            self.trainingFile = flags.train[0]  # only take first entry
            self.training = True
            print('Training flag detected; Entering Training State')
        else:
            self.training = False
        if flags.debug is not False:
            self.debug = True
            print('Debug State: True')

    def train(self):
        filePath = os.path.join(os.getcwd(), self.trainingFile)
        print('\nFile Path:', filePath)
        df = pd.read_csv(filePath, header=0)
        names = df[df.columns[0]]
        numMembers = np.ones(4)
        num = int(ceil(len(names)/4+.5))  # integer division
        numMembers = numMembers*num
        print('Total number of members:', len(names))
        print('Total number of members:', len(names)//4)
        print('Number of members in each house (before remainder)', num)
        print('Remainder:', len(names) % 4)
        for i in range(0, len(names) % 4):  # randomly select a house to award the remainder members to
            j = random.randrange(0, 4)
            numMembers[j] = numMembers[j] + 1
        print('\nHere are the resulting house member counts:\n')
        print(self.houses)
        print(numMembers)

        numbersDist = []  # empty array to hold random numbers generated below
        for i in range(0, len(names)):
            numbers.append(random.random())

    def set_Constants(self):
        self.houses = ['Slytherin', 'Gryffindor', 'Hufflepuff', 'Ravenclaw']
        self.specialNames = {
            'harry potter': 'Gryffindor',
            'hermione granger': 'Gryffindor',
            'ron weasley': 'Gryffindor',
            'draco malfoy': 'Slytherin',
            'lord voldemort': 'Slytherin',
            'albus dumbledore': 'Gryffindor',
            'cameron flannery': 'Gryffindor',
            'michael phalen': 'Gryffindor',
            'alex lim': 'Hufflepuff',
            'parker brown': 'Ravenclaw',
            'brandon nyugen': 'Slytherin',
        }
        self.nameSpecifiedQuotes = {
            'harry potter': 0,
            'ron weasley': 1,
        }
        self.quotes = ['Hmm, difficult. VERY difficult. Plenty of courage, I see. Not a bad mind, either. There\'s talent, oh yes. And a thirst to prove yourself. But where to put you?',
                       'Ah! I know just what to do with you...',
                       'I\'ll have alook inside your mind and tell where you belong!',
                       'Curious, very curious...']

    def get_name(self):
        while True:
            if self.debug:
                break
            self.rawName = input('Enter your first and last name (ex: \'Nikola Tesla\'): ')
            if self.isEmpty(self.rawName):
                print('Error: Please enter your name.')
                continue
            else:
                break

        self.rawName = self.rawName.lower()

    def sort(self):
        if not(self.isPM):
            m=re.sub("([a-zA-Z])[a-zA-Z]*?[^a-zA-Z]*?([a-zA-Z]+?)[^a-zA-Z]*?$", r"\1\2", self.rawName)
            m=re.sub("[^a-zA-Z]","",m)
            hasher = hashlib.sha256()
            hasher.update(m.encode('utf-8'))
            random.seed(hasher.hexdigest(),version=2)
            for j in range(0,4):
                hasher.update(repr(random.random()).encode('utf-8'))
            self.houseChoice = int(hasher.hexdigest(), 16) % 4
            self.sortedHouse = self.houses[self.houseChoice]
        if not(self.debug):
            time.sleep(0.5)
            if self.quoteSpecified:
                self.get_quote(self.nameSpecifiedQuotes[self.rawName])
            else:
                self.get_quote()
            time.sleep(1.5)

    def show_result(self):
        houseText = AsciiText(self.sortedHouse.upper()+'!')
        houseText.display_text()
        print('\n')

    def print_crest(self):
        """print house crest *** NOT CURRENTLY USED*** """
        # prints ascii art of rocket stored in /resources/openrocketengine
        path = os.path.dirname(__file__)
        fname = os.path.join(path, 'assets', self.sortedHouse+'.jpg')

        logo_image = AsciiImage(fname)
        logo_image.display_image()

    def check_name(self):
        """check_name checks to see if the entered name is a principal member.
        If yes, house is predetermined with dictionary"""

        try:
            self.sortedHouse = self.specialNames[self.rawName]
            self.isPM = True
        except KeyError:
            self.isPM = False  # bool: is principal member of club
        try:
            self.nameSpecifiedQuotes[self.rawName]
            self.quoteSpecified = True
        except KeyError:
            self.quoteSpecified = False  # bool: is principal member of club

    def get_quote(self, *arg):
        if arg:
            i = arg[0]
        else:
            i = random.randrange(0, len(self.quotes), 1)
        self.delay_print(i)

    def delay_print(self, i):
        delayTime = 0.05
        for c in self.quotes[i]:
            if self.is_punct_char(c):
                delayTime = 0.15
            else:
                delayTime = 0.05
            sys.stdout.write('%s' % c)
            sys.stdout.flush()
            time.sleep(delayTime)

    def is_punct_char(self, c):
        '''check if char is punctuation char'''
        if c in string.punctuation:
            return 1
        else:
            return 0

    def isEmpty(self, text):
        if text:
            return False
        else:
            return True

if __name__ == '__main__':
    cond = True
    while cond:
        try:
            sp.call('clear')
        except OSError:
            sp.call('cls', shell=True)
        try:
            RPLsorting = sortingHat()
            input("Press Enter to Continue")
        except KeyboardInterrupt:
            print('\nKeyboardInterrupt: Exiting')
            cond = False


import tf
import os
import sys
import time
import copy
import rospy
import ctypes
import struct
import get_info
import numpy as np
from model import *
import open3d as o3d
import indoor3d_util
from natsort import natsorted

from sensor_msgs.msg import PointField
import message_filters
import std_msgs.msg
import sensor_msgs.point_cloud2 as pc2
from sensor_msgs.msg import PointCloud2

class Pointcloud_Sub:

    def __init__(self, name):
        self.name = name

        # init info
        self.info_world_all = list()

        # Params inference
        self.fps = 4.0                # target fps        //PARAM
        self.period = 1.0/self.fps    # target period     //PARAM

        self.init = False
        self.new_pc = False

        # listener
        self.listener = tf.TransformListener()

        # set subscribers
        pc_info = message_filters.Subscriber('/stereo_down/scaled_x2/points2_info', PointCloud2)     # //PARAM
        pc_info.registerCallback(self.cb_pc)

        # set publishers
        self.pub_pc_info_world = rospy.Publisher("/stereo_down/scaled_x2/points2_info_world", PointCloud2, queue_size=4)

        # Set segmentation timer
        rospy.Timer(rospy.Duration(self.period), self.run)

    def cb_pc(self, pc):
        self.pc = pc
        self.new_pc = True

    def run(self,_):
        rospy.loginfo('[%s]: Running', self.name)
        t0 = rospy.Time.now()

        # New pc available
        if not self.new_pc:
            rospy.loginfo('[%s]: No new pointcloud', self.name)
            return
        self.new_pc = False

        # Retrieve image
        try:
            pc = self.pc
            header = self.pc.header
            if not self.init:
                rospy.loginfo('[%s]: Start frame_id change', self.name)
        except:
            rospy.logwarn('[%s]: There is no input pc to run the frame_id change', self.name)
            return

        left_frame_id = "turbot/stereo_down/left_optical"
        world_frame_id = "world_ned"
        left2worldned = self.get_transform(world_frame_id, left_frame_id)

        header.frame_id = "world_ned"

        pc_np_info = self.pc_info2array(pc)
        pc_np_info_world = pc_np.copy()

        if isinstance(left2worldned,int) == False:

            for i in range(pc_np.shape[0]):
                xyz = np.array([[pc_np_info[i,0]],
                                 [pc_np_info[i,1]],
                                 [pc_np_info[i,2]],
                                 [1]])
                xyz_trans_rot = np.matmul(left2worldned, xyz)
                pc_np_info_world[i,0:3] = [xyz_trans_rot[0], xyz_trans_rot[1], xyz_trans_rot[2]]

        # descifrar pc_np_info_world into pipe_list, valve_list, conn_list

        # juntar (si se juntan pipes, pasar union de puntos por 128)
        # si se juntan pipes, recarcular connexions y pipes_near

        # si k = x, hacer check de info world all y borrar instances aparecidas menos de y veces, k = 0
        # si se borra pipe, recalcular pipes near

        if len(info3)>0:
            pc_np_info_world_all = get_info.info_to_array(self.info_world_all)
            pc_info_world_all = self.array2pc_info(header, pc_np_info_world_all)
            self.pub_pc_info_world.publish(pc_info_world_all)

    def get_transform(self, parent, child):
        try:
            rospy.logwarn("[%s]: waiting transform from %s to %s", self.name, parent, child)
            self.listener.waitForTransform(parent, child, rospy.Time(), rospy.Duration(0.3))
            (trans, rot) = self.listener.lookupTransform(parent, child, rospy.Time())
            rospy.loginfo("[%s]: transform for %s found", self.name, child)
            transform = tf.transformations.concatenate_matrices(tf.transformations.translation_matrix(trans), tf.transformations.quaternion_matrix(rot))
        except (tf.LookupException, tf.ConnectivityException, tf.ExtrapolationException, tf.Exception):
            rospy.logerr('[%s]: define %s transform!', self.name, child)
            transform = 0
        return transform

    def pc_info2array(self, ros_pc):
        gen = pc2.read_points(ros_pc, skip_nans=True)   # ROS pointcloud into generator
        pc_np = np.array(list(gen))                     # generator to list to numpy
        pc_np = np.delete(pc_np, 3, 1)
        return pc_np

    def array2pc_info(self, header, array):

        fields =   [PointField('x', 0, PointField.FLOAT32, 1),
                    PointField('y', 4, PointField.FLOAT32, 1),
                    PointField('z', 8, PointField.FLOAT32, 1),
                    PointField('rgba', 12, PointField.UINT32, 1),
                    PointField('t', 16, PointField.FLOAT32, 1),
                    PointField('info', 20, PointField.FLOAT32, 1),
                    PointField('c', 24, PointField.FLOAT32, 1),
                    PointField('inst', 28, PointField.FLOAT32, 1)]

        points = list()

        for i, p in enumerate(array):
            r = int(p[3])
            g = int(p[4])
            b = int(p[5])
            a = 255
            rgb = struct.unpack('I', struct.pack('BBBB', b, g, r, a))[0]

            p_rgb = [p[0], p[1], p[2], rgb, p[6], p[7], p[8], p[9]]
            points.append(p_rgb)

        pc = pc2.create_cloud(header, fields, points)
        return pc

if __name__ == '__main__':
    try:
        rospy.init_node('sub_pc2')
        Pointcloud_Sub(rospy.get_name())

        rospy.spin()
    except rospy.ROSInterruptException:
        pass


import math
import numpy as np
try:
    from tcpb.tcpb import TCProtobufClient
except ImportError:
    pass
import os
import errno

#################################################
### electronic structure routines go here #######
#################################################

#each electronic structure method requires at least two routines:
#1) compute_elec_struct_, which computes energies, forces, and wfs
#2) init_h5_datasets_, which defines the datasets to be output to hdf5
#3) potential_specific_traj_copy, which copies data that is potential specific
#   from one traj data structure to another.  This is used when new trajectories
#   and centroids are spawned.
#other ancillary routines may be included as well

### terachem_cas electronic structure ###
def compute_elec_struct(self,zbackprop):
    if not zbackprop:
        cbackprop = ""
    else:
        cbackprop = "backprop_"

    istate = self.get_istate()
    nstates = self.get_numstates()

    # initialize electronic_phases if not present
    if not hasattr(self,'electronic_phases'):
        self.electronic_phases = np.ones(nstates)
    if not hasattr(self,'backprop_electronic_phases'):
        self.backprop_electronic_phases = np.ones(nstates)

    exec("pos = self.get_" + cbackprop + "positions()")
    pos_list = pos.tolist()

    TC = TCProtobufClient(host='localhost', port=54321)

    base_options = self.get_tc_options()

    base_options["castarget"] = istate

    TC.update_options(**base_options)

    TC.connect()

    # Check if the server is available
    avail = TC.is_available()
    #print "TCPB Server available: {}".format(avail)

    # Write CI vectors and orbitals for initial guess and overlaps
    cwd = os.getcwd()
    # Gradient calculation

    # here we call TC once for energies and once for the gradient
    # will eventually be replaced by a more efficient interface
    options = {}
    results = TC.compute_job_sync("energy", pos_list, "bohr", **options)
    #print results

    e = np.zeros(nstates)
    e[0] = results['energy']

    results = TC.compute_job_sync("gradient", pos_list, "bohr", **options)
    #print results

    f = np.zeros((nstates,self.numdims))
    #print "results['gradient'] ", results['gradient']
    #print "results['gradient'].flatten() ", results['gradient'].flatten()
    f[self.istate,:] = -1.0 * results['gradient'].flatten()

    exec("self.set_" + cbackprop + "energies(e)")

    exec("self.set_" + cbackprop + "forces(f)")

def init_h5_datasets(self):
    self.h5_datasets["time"] = 1
    self.h5_datasets["energies"] = self.numstates
    self.h5_datasets["positions"] = self.numdims
    self.h5_datasets["momenta"] = self.numdims
    self.h5_datasets["forces_i"] = self.numdims
    self.h5_datasets_half_step["time_half_step"] = 1

def potential_specific_traj_copy(self,from_traj):
    self.set_tc_options(from_traj.get_tc_options())
    return

def set_tc_options(self, tco):
    self.tc_options = tco.copy()

def get_tc_options(self):
    return self.tc_options.copy()

def get_prev_wf_positions(self):
    return self.prev_wf_positions.copy()

def get_backprop_prev_wf_positions(self):
    return self.backprop_prev_wf_positions.copy()

def get_prev_wf_positions_in_angstrom(self):
    return 0.529177*self.prev_wf_positions

def get_backprop_prev_wf_positions_in_angstrom(self):
    return 0.529177*self.backprop_prev_wf_positions

def set_prev_wf_positions(self,pos):
    self.prev_wf_positions = pos.copy()

def set_backprop_prev_wf_positions(self,pos):
    self.backprop_prev_wf_positions = pos.copy()

def get_electronic_phases(self):
    return self.electronic_phases.copy()

def set_electronic_phases(self, v):
    self.electronic_phases = v.copy()

def get_backprop_electronic_phases(self):
    return self.backprop_electronic_phases.copy()

def set_backprop_electronic_phases(self, v):
    self.backprop_electronic_phases = v.copy()

###end terachem_cas electronic structure section###


""""Run MS²PIP prediction for single peptide."""

import re
import logging

import click
import matplotlib.pyplot as plt
import numpy as np
import xgboost as xgb
import spectrum_utils.plot as sup
import spectrum_utils.spectrum as sus

from ms2pip.config_parser import ConfigParser
from ms2pip.cython_modules import ms2pip_pyx
from ms2pip.exceptions import InvalidModificationFormattingError, InvalidPeptideError
from ms2pip.ms2pipC import MODELS, apply_mods
from ms2pip.peptides import AMINO_ACID_IDS, Modifications, write_amino_accid_masses
from ms2pip.predict_xgboost import validate_requested_xgb_model, initialize_xgb_models

logger = logging.getLogger("ms2pip")

class SinglePrediction:
    """Run MS²PIP prediction for single peptide."""

    def __init__(self, modification_strings=None) -> None:
        """
        Run MS²PIP prediction for single peptide.

        Parameters
        ----------
        modification_strings: list-like
            List of MS²PIP configuration-style modification strings, e.g.
            `Carbamidomethyl,57.02146,opt,C` or `Oxidation,15.994915,opt,M`. See MS²PIP
            README.md for more info.

        Examples
        --------
        >>> from ms2pip.single_prediction import SinglePrediction
        >>> ms2pip_sp = SinglePrediction(
        >>>     modification_strings=[
        >>>         "Carbamidomethyl,57.021464,opt,C"
        >>>     ]
        >>> )
        >>> mz, intensity, annotation = ms2pip_sp.predict(
        >>>     "GSIGECIAEEEEFELDSESNR", "6|Carbamidomethyl", 3
        >>> )

        """
        if not modification_strings:
            modification_strings = []
        self._init_ms2pip(modification_strings)

    def _init_ms2pip(self, modification_strings):
        self.mod_info = Modifications()
        self.mod_info.modifications = {"ptm": {}, "sptm": {}}
        self.mod_info.add_from_ms2pip_modstrings(modification_strings)
        afile = write_amino_accid_masses()
        modfile = self.mod_info.write_modifications_file(mod_type="ptm")
        modfile2 = self.mod_info.write_modifications_file(mod_type="sptm")
        ms2pip_pyx.ms2pip_init(afile, modfile, modfile2)

    def predict(
        self, peptide, modifications, charge, model="HCD", validate_input=True,
    ):
        """
        Predict single peptide spectrum with MS²PIP.

        Parameters
        ----------
        peptide: string
            Unmodified peptide sequence. Only canonical amino acids are allowed, and
            peptide sequence should be of length [3, 100].
        modifications: string
            MS²PIP style-formatted modification string (e.g. `0|Acetyl|5|Oxidation`).
            See MS²PIP README.md for more info.
        charge: int
            Peptide precursor charge
        model: string (default: "HCD")
            MS²PIP model to use, identical to the option in the MS²PIP configuration
            file.
        validate_input: boolean (default: True)
            Whether to validate input peptide and modifications. Disable for speed-up.

        Returns
        -------
        mz: list[float]
            List with fragment ion m/z values in Da.
        intensity: list[float]
            List with TIC-normalized predicted intensities, order matches `mz`
        annotation: list[str]
            List with fragment ion types and series, order matches `mz`

        """
        peptide = peptide.upper().replace("L", "I")

        if validate_input:
            self._validate_sequence(peptide)
            self._validate_mod_string(modifications)

        peptide = np.array(
            [0] + [AMINO_ACID_IDS[x] for x in peptide] + [0], dtype=np.uint16
        )
        modpeptide = apply_mods(peptide, modifications, self.mod_info.ptm_ids)
        model_id = MODELS[model]["id"]
        peaks_version = MODELS[model]["peaks_version"]
        ce = 30

        mz = np.array(ms2pip_pyx.get_mzs(modpeptide, peaks_version))
        if "xgboost_model_files" in MODELS[model].keys():
            validate_requested_xgb_model(
                MODELS[model]["xgboost_model_files"],
                MODELS[model]["model_hash"]
            )
            xgboost_models = initialize_xgb_models(
                MODELS[model]["xgboost_model_files"],
                1
            )
            xgb_vector = np.array(
                ms2pip_pyx.get_vector(peptide, modpeptide, charge),
                dtype=np.uint16
            )
            xgb_vector = xgb.DMatrix(xgb_vector)
            intensity = []
            for ion_type, model_file in xgboost_models.items():
                preds = model_file.predict(xgb_vector)
                if ion_type in ["x", "y", "y2", "z"]:
                    preds = list(np.array(preds[::-1], dtype=np.float32))
                elif ion_type in ["a", "b", "b2", "c"]:
                    preds = list(np.array(preds, dtype=np.float32))
                intensity.append(preds)
            intensity = np.array(intensity)
        else:
            intensity = np.array(ms2pip_pyx.get_predictions(
                peptide, modpeptide, charge, model_id, peaks_version, ce
            ))
        annotation = []
        for ion_type in MODELS[model]["ion_types"]:
            annotation.append(
                [ion_type.lower() + str(i + 1) for i in range(len(mz[MODELS[model]["ion_types"].index(ion_type)]))]
            )
        annotation = np.array(annotation)
        mz = mz.flatten()
        intensity = self._tic_normalize(self._transform(intensity.flatten()))
        annotation = annotation.flatten()
        return mz, intensity, annotation

    def plot_prediction(
        self, peptide, modifications, charge, prediction=None, ax=None, filename=None,
    ):
        """
        Plot MS²PIP-predicted spectrum with spectrum_utils.

        Parameters
        ----------
        peptide: string
            Unmodified peptide sequence. Only canonical amino acids are allowed, and
            peptide sequence should be of length [3, 100].
        modifications: string
            MS²PIP style-formatted modification string (e.g. `0|Acetyl|5|Oxidation`).
            See MS²PIP README.md for more info.
        charge: int
            Peptide precursor charge.
        prediction: tuple or None (default: None)
            Tuple with `ms2pip.single_prediction.SinglePrediction.predict()` output.
        ax: matplotlib.axes.Axes or None (default: None)
            Figure ax to plot onto.
        filename: str or None (default: None)
            Filename to save plot to. File extension defines the format. Figure will
            not be saved if None.

        """
        if not prediction:
            prediction = self.predict(peptide, modifications, charge)
        mz, intensity, annotation = prediction

        identifier = f"{peptide}/{charge}/{modifications}"
        precursor_mz = self.mod_info.calc_precursor_mz(peptide, modifications, charge)
        mod_dict = self._modifications_to_dict(modifications)
        sus_annotation = self._get_sus_annotation(mz, annotation)

        spectrum = sus.MsmsSpectrum(
            identifier,
            precursor_mz,
            charge,
            mz,
            intensity,
            annotation=sus_annotation,
            retention_time=None,
            peptide=peptide,
            modifications=mod_dict,
        )

        if not ax:
            ax = plt.gca()
        sup.spectrum(spectrum, ax=ax)
        ax.set_title("MS²PIP prediction for " + identifier)
        if filename:
            plt.savefig(filename)

    @staticmethod
    def _validate_sequence(sequence):
        """Validate peptide sequence for MS²PIP."""
        pattern = r"[ACDEFGHIKLMNPQRSTVWY]{3,99}"
        if not re.fullmatch(pattern, sequence):
            raise InvalidPeptideError

    @staticmethod
    def _validate_mod_string(mod_string):
        """Validate modification string for MS²PIP."""
        pattern = r"^(?:(?:-1|[0-9]{1,2})\|(?:[^\|\s]+)\|?)*$|^-$"
        if not re.fullmatch(pattern, mod_string):
            raise InvalidModificationFormattingError

    @staticmethod
    def _transform(intensity):
        """Undo MS²PIP peak intensity log transformation and pseudo-count."""
        return (2 ** np.array(intensity)) - 0.001

    @staticmethod
    def _tic_normalize(intensity):
        """TIC-normalize peak intensities."""
        intensity = np.array(intensity)
        return intensity / np.sum(intensity)

    @staticmethod
    def _modifications_to_dict(modifications):
        """Convert ms2pip modification notation to spectrum_utils dict."""
        def parse_loc(loc):
            if loc == "0":
                return "N-term"
            elif loc == "-1":
                return "C-term"
            else:
                return int(loc)

        m_split = [modifications.split("|")[i::2] for i in [0, 1]]
        mods_dict = {parse_loc(loc): name for loc, name in zip(m_split[0], m_split[1])}
        return mods_dict

    @staticmethod
    def _get_sus_annotation(mz, annotation):
        """Get spectrum_utils.PeptideFragmentAnnotation objects."""
        return [
            sus.PeptideFragmentAnnotation(1, mz, annotation[0], annotation[1:])
            for mz, annotation in zip(mz, annotation)
        ]

@click.command()
@click.argument("peptide", type=str)
@click.argument("modifications", type=str)
@click.argument("charge", type=int)
@click.option("-m", "--model", type=str, default="HCD", help="")
@click.option("-c", "--configfile", type=click.Path(exists=True), default=None, help="")
@click.option(
    "-o",
    "--output",
    type=click.Path(exists=False),
    default="ms2pip_prediction.png",
    help="",
)
def _main(
    peptide,
    modifications,
    charge,
    model="HCD",
    configfile=None,
    output="ms2pip_prediction.png",
):
    """
    Generate MS²PIP-predicted spectrum and plot.

    \b
    Examples:
     - ms2pip-single-prediction PGAQANPYSR "-" 3
     - ms2pip-single-prediction -o prediction.png PGAQANPYSR "-" 3
     - ms2pip-single-prediction -c config.toml NSVPCSR "5|Carbamidomethyl" 3

    """
    root_logger = logging.getLogger()
    handler = logging.StreamHandler()
    root_logger.addHandler(handler)
    root_logger.setLevel(logging.INFO)

    if configfile:
        config_parser = ConfigParser(configfile)
        mod_strings = config_parser.config["ms2pip"]["ptm"]
    else:
        mod_strings = None
    ms2pip_sp = SinglePrediction(mod_strings)

    _, ax = plt.subplots(figsize=(10, 5))
    prediction = ms2pip_sp.predict(peptide, modifications, charge, model=model)
    ms2pip_sp.plot_prediction(
        peptide, modifications, charge, prediction=prediction, ax=ax, filename=output,
    )

if __name__ == "__main__":
    _main()


# Taken from https://github.com/bastibe/SoundCard
import soundcard as sc
import numpy

# get a list of all speakers:
speakers = sc.all_speakers()
# get the current default speaker on your system:
default_speaker = sc.default_speaker()
# get a list of all microphones:
mics = sc.all_microphones()
# get the current default microphone on your system:
default_mic = sc.default_microphone()

# record and play back one second of audio:
# data = default_mic.record(samplerate=48000, numframes=48000)
# default_speaker.play(data/numpy.max(data), samplerate=48000)

# alternatively, get a `Recorder` and `Player` object
# and play or record continuously:
while True:
    with default_mic.recorder(samplerate=48000) as mic, \
        default_speaker.player(samplerate=48000) as sp:
        for _ in range(100):
            data = mic.record(numframes=1024)
            sp.play(data)