import gzip
import json
import logging
import os
import re
from collections import abc
from contextlib import contextmanager

import math
import matplotlib
import sys
from scipy import signal

matplotlib.use("Qt5Agg")
os.environ['QT_API'] = 'pyqt5'
# os.environ['PYQTGRAPH_QT_LIB'] = 'PyQt5'

from model.waveform import WaveformController
from model.checker import VersionChecker, ReleaseNotesDialog
from model.report import SaveReportDialog, block_signals
from model.preferences import DISPLAY_SHOW_FILTERED_SIGNALS, SYSTEM_CHECK_FOR_UPDATES, BIQUAD_EXPORT_MAX, \
    BIQUAD_EXPORT_FS, BIQUAD_EXPORT_DEVICE, SHOW_NO_FILTERS, SYSTEM_CHECK_FOR_BETA_UPDATES, \
    BEQ_REPOS, BEQ_DEFAULT_REPO

from ui.delegates import RegexValidator

import pyqtgraph as pg
import qtawesome as qta
from matplotlib import style

from model.iir import Passthrough, CompleteFilter, as_equalizer_apo, COMBINED
from ui.biquad import Ui_exportBiquadDialog
from ui.savechart import Ui_saveChartDialog

from qtpy import QtCore, QtWidgets
from qtpy.QtCore import QSettings, QThreadPool, QUrl
from qtpy.QtGui import QIcon, QFont, QCursor, QTextCursor, QDesktopServices
from qtpy.QtWidgets import QMainWindow, QApplication, QErrorMessage, QAbstractItemView, QDialog, QFileDialog, \
    QHeaderView, QMessageBox, QHBoxLayout, QToolButton

from model.preferences import PreferencesDialog, BINARIES_GROUP, ANALYSIS_TARGET_FS, STYLE_MATPLOTLIB_THEME, \
    Preferences, STYLE_MATPLOTLIB_THEME_DEFAULT, SCREEN_GEOMETRY, SCREEN_WINDOW_STATE, FILTERS_PRESET_x, \
    DISPLAY_SHOW_LEGEND, DISPLAY_SHOW_FILTERS, SHOW_FILTER_OPTIONS, SHOW_SIGNAL_OPTIONS, DISPLAY_SHOW_SIGNALS, \
    SHOW_FILTERED_SIGNAL_OPTIONS
from ui.beq import Ui_MainWindow

logger = logging.getLogger('beq')
logging.getLogger('matplotlib').setLevel(logging.WARNING)

@contextmanager
def wait_cursor(msg=None):
    '''
    Allows long running functions to show a busy cursor.
    :param msg: a message to put in the status bar.
    '''
    try:
        QApplication.setOverrideCursor(QCursor(QtCore.Qt.WaitCursor))
        yield
    finally:
        QApplication.restoreOverrideCursor()

class BeqDesigner(QMainWindow, Ui_MainWindow):
    '''
    The main UI.
    '''

    def __init__(self, app, prefs, parent=None):
        super(BeqDesigner, self).__init__(parent)
        self.logger = logging.getLogger('beqdesigner')
        self.app = app
        self.preferences = prefs
        if getattr(sys, 'frozen', False):
            self.__style_path_root = sys._MEIPASS
        else:
            self.__style_path_root = os.path.dirname(__file__)
        self.__version = '0.0.0-alpha.1'
        v_path = os.path.abspath(os.path.join(self.__style_path_root, 'VERSION'))
        try:
            with open(v_path) as version_file:
                self.__version = version_file.read().strip()
        except:
            logger.exception(f"Unable to read {v_path}")
        if self.preferences.get(SYSTEM_CHECK_FOR_UPDATES):
            QThreadPool.globalInstance().start(VersionChecker(self.preferences.get(SYSTEM_CHECK_FOR_BETA_UPDATES),
                                                              self.__alert_on_old_version,
                                                              self.__alert_on_version_check_fail,
                                                              self.__version))

        matplotlib_theme = self.preferences.get(STYLE_MATPLOTLIB_THEME)
        if matplotlib_theme is not None:
            if matplotlib_theme.startswith('beq'):
                style.use(os.path.join(self.__style_path_root, 'style', 'mpl', f"{matplotlib_theme}.mplstyle"))
            else:
                style.use(matplotlib_theme)

        pg.setConfigOption('background', matplotlib.colors.to_hex(matplotlib.rcParams['axes.facecolor']))
        pg.setConfigOption('foreground', matplotlib.colors.to_hex(matplotlib.rcParams['axes.edgecolor']))
        pg.setConfigOption('leftButtonPan', False)
        self.setupUi(self)
        self.__decorate_splitter()
        self.limitsButton.setIcon(qta.icon('fa5s.arrows-alt'))
        self.showValuesButton.setIcon(qta.icon('ei.eye-open'))
        # init beq repos for 0.9.0 backwards compatibility
        self.__ensure_beq_repos_exist()
        # logs
        from model.log import RollingLogger
        self.logViewer = RollingLogger(self.preferences, parent=self)
        self.actionShow_Logs.triggered.connect(self.logViewer.show_logs)
        self.actionPreferences.triggered.connect(self.showPreferences)
        # init a default signal for when we want to edit a filter without a signal
        from model.signal import SingleChannelSignalData, Signal
        default_fs = self.preferences.get(ANALYSIS_TARGET_FS)
        default_signal = Signal('default', signal.unit_impulse(default_fs, 'mid'), self.preferences, fs=default_fs)
        self.__default_signal = SingleChannelSignalData(name='default',
                                                        signal=default_signal,
                                                        filter=CompleteFilter(fs=default_fs))
        # init the filter view selector
        self.showFilters.blockSignals(True)
        for x in SHOW_FILTER_OPTIONS:
            self.showFilters.addItem(x)
        selected = self.preferences.get(DISPLAY_SHOW_FILTERS)
        selected_idx = self.showFilters.findText(selected)
        if selected_idx != -1:
            self.showFilters.setCurrentIndex(selected_idx)
        else:
            logger.info(f"Ignoring unknown cached preference for {DISPLAY_SHOW_FILTERS} - {selected}")
        self.showFilters.blockSignals(False)
        # filter view/model
        self.actionShow_Filter_Widget.triggered.connect(lambda:self.editFilter(small=True))
        self.filterView.setSelectionBehavior(QAbstractItemView.SelectRows)
        self.filterView.doubleClicked.connect(self.editFilter)
        from model.filter import FilterTableModel, FilterModel
        self.__filter_model = FilterModel(self.filterView, self.preferences, label=self.filtersLabel,
                                          on_update=self.on_filter_change)
        self.__filter_table_model = FilterTableModel(self.__filter_model, parent=parent)
        self.filterView.setModel(self.__filter_table_model)
        self.filterView.selectionModel().selectionChanged.connect(self.on_filter_selected)
        self.filterView.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        for i in range(1, 4):
            getattr(self, f"action_load_preset_{i}").triggered.connect(self.load_preset(i))
            getattr(self, f"action_clear_preset_{i}").triggered.connect(self.clear_preset(i))
            getattr(self, f"action_store_preset_{i}").triggered.connect(self.set_preset(i))
            self.enable_preset(i)
        self.actionAdd_BEQ_Filter.triggered.connect(self.add_beq_filter)
        self.actionClear_Filters.triggered.connect(self.clearFilters)
        # init the signal view selector
        self.showSignals.blockSignals(True)
        for x in SHOW_SIGNAL_OPTIONS:
            self.showSignals.addItem(x)
        selected = self.preferences.get(DISPLAY_SHOW_SIGNALS)
        selected_idx = self.showSignals.findText(selected)
        if selected_idx != -1:
            self.showSignals.setCurrentIndex(selected_idx)
        else:
            logger.info(f"Ignoring unknown cached preference for {DISPLAY_SHOW_SIGNALS} - {selected}")
        self.showSignals.blockSignals(False)
        # init the filtered signal view selecter
        self.showFilteredSignals.blockSignals(True)
        for x in SHOW_FILTERED_SIGNAL_OPTIONS:
            self.showFilteredSignals.addItem(x)
        selected = self.preferences.get(DISPLAY_SHOW_FILTERED_SIGNALS)
        selected_idx = self.showFilteredSignals.findText(selected)
        if selected_idx != -1:
            self.showFilteredSignals.setCurrentIndex(selected_idx)
        else:
            logger.info(f"Ignoring unknown cached preference for {DISPLAY_SHOW_FILTERED_SIGNALS} - {selected}")
        self.showFilteredSignals.blockSignals(False)
        # signal model
        self.__cached_selected_signal = None
        self.__configure_signal_model(parent)
        # signal editor
        self.actionMerge_Signals.triggered.connect(self.__open_merge_signal_dialog)
        self.actionMerge_Signals.setEnabled(False)
        # magnitude
        self.showLegend.setChecked(bool(self.preferences.get(DISPLAY_SHOW_LEGEND)))
        from model.magnitude import MagnitudeModel
        self.__magnitude_model = MagnitudeModel('main', self.mainChart, self.preferences,
                                                self.__signal_model.get_curve_data, 'Signals',
                                                secondary_data_provider=self.__filter_model.get_curve_data,
                                                secondary_name='Filters',
                                                show_legend=lambda: self.showLegend.isChecked(), allow_line_resize=True)
        self.__filter_model.filter = self.__default_signal.filter
        self.y1MaxPlus10Button.setIcon(qta.icon('mdi.chevron-double-up'))
        self.y1MaxPlus5Button.setIcon(qta.icon('mdi.chevron-up'))
        self.y1MaxMinus5Button.setIcon(qta.icon('mdi.chevron-down'))
        self.y1MaxMinus10Button.setIcon(qta.icon('mdi.chevron-double-down'))
        self.y1MinPlus10Button.setIcon(qta.icon('mdi.chevron-double-up'))
        self.y1MinPlus5Button.setIcon(qta.icon('mdi.chevron-up'))
        self.y1MinMinus5Button.setIcon(qta.icon('mdi.chevron-down'))
        self.y1MinMinus10Button.setIcon(qta.icon('mdi.chevron-double-down'))
        self.y1AutoOnButton.setIcon(qta.icon('fa5s.magic'))
        self.y1AutoOnButton.setChecked(True)
        self.y1AutoOffButton.setIcon(qta.icon('fa5s.magic', 'fa5s.ban', options=[{'scale_factor': 0.75}, {}]))
        self.y2MaxPlus10Button.setIcon(qta.icon('mdi.chevron-double-up'))
        self.y2MaxPlus5Button.setIcon(qta.icon('mdi.chevron-up'))
        self.y2MaxMinus5Button.setIcon(qta.icon('mdi.chevron-down'))
        self.y2MaxMinus10Button.setIcon(qta.icon('mdi.chevron-double-down'))
        self.y2MinPlus10Button.setIcon(qta.icon('mdi.chevron-double-up'))
        self.y2MinPlus5Button.setIcon(qta.icon('mdi.chevron-up'))
        self.y2MinMinus5Button.setIcon(qta.icon('mdi.chevron-down'))
        self.y2MinMinus10Button.setIcon(qta.icon('mdi.chevron-double-down'))
        self.y2AutoOnButton.setIcon(qta.icon('fa5s.magic'))
        self.y2AutoOnButton.setChecked(True)
        self.y2AutoOffButton.setIcon(qta.icon('fa5s.magic', 'fa5s.ban', options=[{'scale_factor': 0.75}, {}]))
        # waveform
        self.__waveform_controller = WaveformController(self.preferences,
                                                        self.__signal_model,
                                                        self.waveformChart,
                                                        self.filteredSpectrumChart,
                                                        self.signalSelector,
                                                        self.rmsLevel,
                                                        self.headroom,
                                                        self.crestFactor,
                                                        self.bmHeadroom,
                                                        self.bmlpfPosition,
                                                        self.bmClipBefore,
                                                        self.bmClipAfter,
                                                        self.waveformIsFiltered,
                                                        self.hardClipWaveform,
                                                        self.startTime,
                                                        self.endTime,
                                                        self.showSpectrumButton,
                                                        self.hideSpectrumButton,
                                                        self.zoomInButton,
                                                        self.zoomOutButton,
                                                        self.compareSpectrumButton,
                                                        self.sourceFile,
                                                        self.loadSignalButton,
                                                        self.filteredSpectrumLimitsButton,
                                                        self.showStatsButton,
                                                        self.yMin,
                                                        self.yMax,
                                                        self.bmhpfOn,
                                                        self.saveWaveformChartButton)
        self.__hide_waveform_chart()
        self.actionClear_Signals.triggered.connect(self.clearSignals)
        # processing
        self.ensurePathContainsExternalTools()
        # extraction
        self.actionExtract_Audio.triggered.connect(self.showExtractAudioDialog)
        self.action_Remux_Audio.triggered.connect(self.showRemuxAudioDialog)
        self.action_Batch_Extract.triggered.connect(self.showBatchExtractDialog)
        # analysis
        self.actionAnalyse_Audio.triggered.connect(self.showAnalyseAudioDialog)
        # import
        self.actionLoad_Filter.triggered.connect(self.importFilter)
        self.actionLoad_Signal.triggered.connect(self.importSignal)
        self.action_Load_Project.triggered.connect(self.importProject)
        # export
        self.actionSave_Report.triggered.connect(self.exportReport)
        self.actionSave_Chart.triggered.connect(self.exportChart)
        self.actionExport_Biquad.triggered.connect(self.exportBiquads)
        self.actionSave_Filter.triggered.connect(self.exportFilter)
        self.actionExport_FRD.triggered.connect(self.showExportFRDDialog)
        self.actionSave_Signal.triggered.connect(self.showExportSignalDialog)
        self.action_Save_Project.triggered.connect(self.exportProject)
        self.actionAbout.triggered.connect(self.showAbout)
        # tools
        self.actionMerge_Minidsp_XML.triggered.connect(self.merge_minidsp_xml)
        self.actionCreate_AVS_Post.triggered.connect(self.create_avs_post)
        self.actionUser_Guide.triggered.connect(self.show_help)
        self.actionRelease_Notes.triggered.connect(self.show_release_notes)
        self.actionBrowse_Catalogue.triggered.connect(self.show_catalogue)
        self.actionExport_BEQ_Filter.triggered.connect(self.export_beq_filter)
        self.actionSync_with_HTP_1.triggered.connect(self.sync_htp1)
        self.actionManage_JRiver_MC.triggered.connect(self.__import_jriver_peq)

    def __ensure_beq_repos_exist(self):
        '''
        Allows the user to choose whether to add MOberhardt's BEQ repo.
        This should only ever be displayed once on upgrade beyond 0.9.0
        '''
        if not self.preferences.has(BEQ_REPOS):
            repo_list = [BEQ_DEFAULT_REPO]
            result = QMessageBox.question(self,
                                          'Add MOberhardt BEQ Repo?',
                                          f"Do you want to use MOberhardt's BEQs?",
                                          QMessageBox.Yes | QMessageBox.No,
                                          QMessageBox.No)
            if result == QMessageBox.Yes:
                repo_list.append('https://github.com/Mobe1969/miniDSPBEQ.git')
            self.preferences.set(BEQ_REPOS, '|'.join(repo_list))

    def show_release_notes(self):
        ''' Shows the release notes '''
        QThreadPool.globalInstance().start(VersionChecker(self.preferences.get(SYSTEM_CHECK_FOR_BETA_UPDATES),
                                                          self.__alert_on_old_version,
                                                          self.__alert_on_version_check_fail,
                                                          self.__version,
                                                          signal_anyway=True))

    def show_help(self):
        ''' Opens the user guide in a browser '''
        QDesktopServices.openUrl(QUrl('https://beqdesigner.readthedocs.io/en/latest'))

    def smoothSignals(self):
        ''' Applies the current smoothing options to the visible signals. '''
        fraction_idx = self.octaveSmoothing.currentIndex()
        if fraction_idx == 0:
            smooth_type = 0
        else:
            try:
                smooth_type = int(self.octaveSmoothing.currentText()[2:])
            except:
                smooth_type = 'SG'
        with wait_cursor():
            if self.smoothAllSignals.isChecked():
                for s in self.__signal_model:
                    s.smooth(smooth_type)
            else:
                signal_select = self.signalView.selectionModel()
                if signal_select.hasSelection() and len(signal_select.selectedRows()) == 1:
                    signal_data = self.__signal_model[signal_select.selectedRows()[0].row()]
                    if signal_data.signal is not None:
                        signal_data.smooth(smooth_type)
        self.__magnitude_model.redraw()
        self.signalView.viewport().update()

    def __decorate_splitter(self):
        '''
        Adds some visible handles to the splitter so we can quickly show/hide it.
        '''
        handle = self.chartSplitter.handle(1)
        self.chartSplitter.setHandleWidth(15)
        layout = QHBoxLayout()
        layout.setContentsMargins(0, 0, 0, 0)
        button = QToolButton(handle)
        button.setIcon(qta.icon('fa5s.chevron-up'))
        button.clicked.connect(self.__show_waveform_chart)
        layout.addWidget(button)
        button = QToolButton(handle)
        button.setIcon(qta.icon('fa5s.chevron-down'))
        button.clicked.connect(self.__hide_waveform_chart)
        layout.addWidget(button)
        handle.setLayout(layout)

    def __hide_waveform_chart(self):
        self.chartSplitter.setSizes([1, 0])
        self.__waveform_controller.on_visibility_change(show=False)

    def __show_waveform_chart(self):
        self.chartSplitter.setSizes([100000, 100000])
        self.__waveform_controller.on_visibility_change(show=True)

    def __alert_on_version_check_fail(self, message):
        '''
        Displays an alert if the version check fails.
        :param message: the message.
        '''
        msg_box = QMessageBox()
        msg_box.setText(message)
        msg_box.setIcon(QMessageBox.Warning)
        msg_box.setWindowTitle('Unable to Complete Version Check')
        msg_box.exec()

    def __alert_on_old_version(self, versions, issues):
        ''' Presents a dialog if there is a new version available. '''
        ReleaseNotesDialog(self, versions, issues).exec()

    def __configure_signal_model(self, parent):
        '''
        Wires up the signal model so that
        * only one row can be selected at a time
        * there is always a row selected
        * the filter model is always honouring the selected row correctly (where correctly means the filter model is
        operating on the filter of the selected signal if it is an unlinked signal or a master OR the filter of the
        master of the selected signal if the signal is a slave).
        This last point is crucial as otherwise the linked signals won't update at the same time.
        '''
        self.signalView.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        self.signalView.setSelectionBehavior(QAbstractItemView.SelectRows)
        self.signalView.setSelectionMode(QAbstractItemView.SingleSelection)
        from model.signal import SignalModel, SignalTableModel
        self.__signal_model = SignalModel(self.signalView, self.__default_signal, self.preferences,
                                          on_update=self.on_signal_change)
        self.__signal_table_model = SignalTableModel(self.__signal_model, parent=parent)
        self.signalView.setModel(self.__signal_table_model)
        self.signalView.selectionModel().selectionChanged.connect(self.on_signal_selected)
        self.signalView.model().dataChanged.connect(self.on_signal_data_change)
        self.signalView.model().modelAboutToBeReset.connect(self.on_signal_data_about_to_reset)
        self.signalView.model().modelReset.connect(self.on_signal_data_reset)
        self.signalView.model().rowsInserted.connect(self.__handle_signals_inserted)
        self.signalView.model().rowsRemoved.connect(self.__handle_signals_removed)
        self.signalView.setItemDelegateForColumn(0, RegexValidator('^.+$'))

    def on_signal_data_about_to_reset(self):
        '''
        Caches the selected signal so we can restore it post reset.
        '''
        signal_select = self.signalView.selectionModel()
        if signal_select.hasSelection() and len(signal_select.selectedRows()) == 1:
            self.__cached_selected_signal = self.__signal_model[signal_select.selectedRows()[0].row()]

    def on_signal_data_reset(self):
        '''
        Restores the cached selected signal if we have one.
        '''
        selected_row = -1
        if self.__cached_selected_signal is not None:
            for idx, s in enumerate(self.__signal_model):
                if s.name == self.__cached_selected_signal.name:
                    selected_row = idx
            self.__cached_selected_signal = None
        else:
            signal_count = len(self.__signal_model)
            if signal_count > 0:
                selected_row = signal_count - 1
        if selected_row > -1:
            logger.debug(f"Selected row {selected_row} after signal data reset")
            self.signalView.selectRow(selected_row)

    def on_signal_data_change(self):
        '''
        Delegates to on_signal_change which we can't call directly because of the argument mismatch (dataChanged emits
        the QModelIndexes specifying which data has changed which we don't care about).
        '''
        logger.debug(f"Handling signal data change")
        self.on_signal_selected()
        self.on_signal_change()

    def on_signal_change(self, names=None):
        '''
        Reacts to a change in the signal model by updating the reference and redrawing the chart.
        :param names: the signal names.
        '''
        if names is not None:
            self.update_reference_series(names, self.signalReference, True)
        self.__waveform_controller.refresh_selector()
        if self.signalSelector.count() == 1:
            self.__hide_waveform_chart()
        self.linkSignalButton.setEnabled(len(self.__signal_model) > 1)
        self.__magnitude_model.redraw()
        signal_count = sum(s.signal is not None for s in self.__signal_model.non_bm_signals)
        signal_count += sum([len(bm.channels) for bm in self.__signal_model.bass_managed_signals])
        self.actionMerge_Signals.setEnabled(signal_count > 1)

    def __get_selected_signal(self):
        '''
        :return: the selected signal (as selected in the table going back to its master if the selected signal is a
        slave) or the default signal if nothing is selected.
        '''
        signal_select = self.signalView.selectionModel()
        if signal_select.hasSelection() and len(signal_select.selectedRows()) == 1:
            signal = self.__signal_model[signal_select.selectedRows()[0].row()]
            if signal.master is not None:
                signal = signal.master
        else:
            signal = self.__signal_model.default_signal
        return signal

    def __handle_signals_inserted(self, idx, first, last):
        '''
        Selects the last signal added by default.
        :param idx: the idx.
        :param first: the first row inserted.
        :param last: the last row inserted.
        '''
        logger.debug(f"Selecting signal {last} on insert")
        self.signalView.selectRow(last)

    def __handle_signals_removed(self, idx, first, last):
        '''
        Ensures the correct signal is selected when signals are removed. This defaults to the last signal in the table
        if there are any signals otherwise it reverts to the default signal.
        :param idx: the idx.
        :param first: first row removed.
        :param last: last row removed.
        '''
        signal_count = len(self.__signal_model)
        if signal_count > 0:
            logger.debug(f"Selecting signal {signal_count - 1} on remove")
            self.signalView.selectRow(signal_count - 1)
        else:
            logger.debug(f"No signals in model, selecting default filter on remove")
            self.on_signal_selected()

    def addSignal(self):
        '''
        Adds signals via the signal dialog.
        '''
        from model.signal import SignalDialog
        SignalDialog(self.preferences, self.__signal_model, parent=self).show()

    def linkSignals(self):
        '''
        Lets the user link signals via a matrix mapping.
        '''
        from model.link import LinkSignalsDialog
        LinkSignalsDialog(self.__signal_model, parent=self).exec()

    def deleteSignal(self):
        '''
        Deletes the currently selected signals.
        '''
        selection = self.signalView.selectionModel()
        if selection.hasSelection():
            self.__signal_model.delete([x.row() for x in selection.selectedRows()])
        if len(self.__signal_model) > 0:
            self.signalView.selectRow(0)
        else:
            self.signalView.clearSelection()
            # nothing in qt appears to emit selectionChanged when you clear the selection so have to call it ourselves
            self.on_signal_selected()

    def clearSignals(self):
        ''' Deletes all signals '''
        while len(self.__signal_model) > 0:
            self.signalView.selectRow(0)
            self.deleteSignal()

    def importSignal(self):
        '''
        Allows the user to load a signal from a saved file.
        '''

        def parser(file_name):
            with gzip.open(file_name, 'r') as infile:
                return json.loads(infile.read().decode('utf-8'))
        input = self.__load('*.signal', 'Load Signal', parser)
        if input is not None:
            with (wait_cursor()):
                from model.codec import signaldata_from_json
                self.__signal_model.add(signaldata_from_json(input, self.preferences))

    def on_signal_selected(self):
        '''
        Enables the edit & delete button if there are selected rows.
        '''
        selection = self.signalView.selectionModel()
        self.clearSignalsButton.setEnabled(len(self.__signal_model) > 0)
        self.deleteSignalButton.setEnabled(selection.hasSelection())
        if len(selection.selectedRows()) == 1:
            signal_data = self.__get_selected_signal()
            self.__filter_model.filter = signal_data.filter
            with block_signals(self.octaveSmoothing):
                if signal_data.smoothing_type is not None:
                    self.octaveSmoothing.setCurrentText(signal_data.smoothing_description)
                else:
                    self.octaveSmoothing.setCurrentText('None')
            self.__waveform_controller.selected_name = self.__signal_model[selection.selectedRows()[0].row()].name
        else:
            if len(self.__filter_model.filter) > 0:
                self.__default_signal.filter = self.__filter_model.filter
            # we have to set the filter on the default signal and then set that back onto the filter model to ensure
            # the right links are established re listeners and label names
            self.__filter_model.filter = self.__default_signal.filter
            with block_signals(self.octaveSmoothing):
                self.octaveSmoothing.setCurrentText('None')

    def on_filter_change(self, names):
        '''
        Reacts to a change in the filter model by updating the reference, redrawing the chart and other filter related
        things.
        :param names: the signal names.
        '''
        self.update_reference_series(names, self.filterReference, False)
        self.__magnitude_model.redraw()
        self.__enable_filter_actions()
        self.__check_active_preset(self.__filter_model.filter.preset_idx)
        self.show_y2_tool_buttons(self.showFilters.currentText() != SHOW_NO_FILTERS and len(self.__filter_model) > 0)

    def importFilter(self):
        '''
        Allows the user to replace the current filter with one loaded from a file.
        '''
        from model.filter import load_filter
        filter_json = load_filter(self, self.statusbar)
        if filter_json is not None:
            self.__apply_filter(filter_json)
            self.__magnitude_model.redraw()

    def __apply_filter(self, filter_json):
        from model.codec import filter_from_json
        signal = self.__get_selected_signal()
        filt = filter_from_json(filter_json).resample(signal.fs)
        signal.filter = filt
        self.__filter_model.filter = filt

    def addFilter(self, small=False):
        '''
        Adds a filter via the filter dialog.
        '''
        from model.filter import FilterDialog
        FilterDialog(self.preferences, self.__get_selected_signal(), self.__filter_model,
                     lambda: self.__magnitude_model.redraw(), parent=self, small=small).show()

    def editFilter(self, small=False):
        '''
        Edits the currently selected filter via the filter dialog.
        '''
        selection = self.filterView.selectionModel()
        if selection.hasSelection() and len(selection.selectedRows()) == 1:
            signal = self.__get_selected_signal()
            from model.filter import FilterDialog
            FilterDialog(self.preferences, signal, self.__filter_model,
                         lambda: self.__magnitude_model.redraw(),
                         selected_filter=signal.filter[selection.selectedRows()[0].row()],
                         parent=self, small=small).show()
        else:
            self.addFilter(small=small)

    def deleteFilter(self):
        '''
        Deletes the selected filters.
        '''
        selection = self.filterView.selectionModel()
        if selection.hasSelection():
            self.__filter_model.delete([x.row() for x in selection.selectedRows()])

    def clearFilters(self):
        ''' Deletes all filters  '''
        while len(self.__filter_model) > 0:
            self.filterView.selectRow(0)
            self.deleteFilter()
        self.on_filter_selected()

    def __enable_filter_actions(self):
        '''
        Enables the save filter if we have filters to save.
        '''
        enabled = len(self.__filter_model) > 0
        self.actionSave_Filter.setEnabled(enabled)
        self.actionClear_Filters.setEnabled(enabled)
        self.clearFiltersButton.setEnabled(enabled)

    def on_filter_selected(self):
        '''
        Enables the edit & delete button if there are selected rows.
        '''
        selection = self.filterView.selectionModel()
        self.deleteFilterButton.setEnabled(selection.hasSelection())
        self.editFilterButton.setEnabled(len(selection.selectedRows()) == 1)

    def exportFilter(self):
        '''
        Allows the user to save the current filter to a file.
        '''
        dialog = QFileDialog(parent=self)
        dialog.setFileMode(QFileDialog.AnyFile)
        dialog.setNameFilter(f"*.filter")
        dialog.setWindowTitle(f"Save Filter")
        dialog.setLabelText(QFileDialog.Accept, 'Save')
        if dialog.exec():
            selected = dialog.selectedFiles()
            if len(selected) > 0:
                if not selected[0].endswith('.filter'):
                    selected[0] += '.filter'
                with open(selected[0], 'w+') as outfile:
                    json.dump(self.__filter_model.filter.to_json(), outfile)
                    self.statusbar.showMessage(f"Saved filter to {outfile.name}")

    def exportChart(self):
        '''
        Saves the currently selected chart to a file.
        '''
        self.__magnitude_model.export_chart(status_bar=self.statusbar)

    def exportReport(self):
        '''
        Saves the currently selected chart to a file.
        '''
        dialog = SaveReportDialog(self, self.preferences, self.__signal_model, self.__filter_model, self.statusbar,
                                  self.__get_selected_signal())
        dialog.exec()

    def exportBiquads(self):
        '''
        Shows the biquads for the current filter set.
        '''
        dialog = ExportBiquadDialog(self.__filter_model.filter, self.preferences)
        dialog.exec()

    def __load(self, filter, title, parser):
        '''
        Presents a file dialog to the user so they can choose something to load.
        :return: the loaded thing, if any.
        '''
        input = None
        dialog = QFileDialog(parent=self)
        dialog.setFileMode(QFileDialog.ExistingFile)
        dialog.setNameFilter(filter)
        dialog.setWindowTitle(title)
        if dialog.exec():
            selected = dialog.selectedFiles()
            if len(selected) > 0:
                input = parser(selected[0])
                if input is not None:
                    self.statusbar.showMessage(f"Loaded {selected[0]}")
        return input

    def ensurePathContainsExternalTools(self):
        '''
        Ensures that all external tool paths are on the path.
        '''
        path = os.environ.get('PATH', [])
        paths = path.split(os.pathsep)
        locs = self.preferences.get_all(BINARIES_GROUP)
        if len(locs) > 0:
            logging.info(f"Adding {locs} to PATH")
            os.environ['PATH'] = os.pathsep.join([l for l in locs if l not in paths]) + os.pathsep + path
        else:
            logger.warning(f"No {BINARIES_GROUP} paths set")
            # TODO attempt to call each binary with a test command to test if they are really on the path

    def setupUi(self, main_window):
        super().setupUi(self)
        geometry = self.preferences.get(SCREEN_GEOMETRY)
        if geometry is not None:
            self.restoreGeometry(geometry)
        else:
            screen_geometry = self.app.desktop().availableGeometry()
            if screen_geometry.height() < 800:
                self.showMaximized()
        window_state = self.preferences.get(SCREEN_WINDOW_STATE)
        if window_state is not None:
            self.restoreState(window_state)

    def closeEvent(self, *args, **kwargs):
        '''
        Saves the window state on close.
        :param args:
        :param kwargs:
        '''
        self.preferences.set(SCREEN_GEOMETRY, self.saveGeometry())
        self.preferences.set(SCREEN_WINDOW_STATE, self.saveState())
        super().closeEvent(*args, **kwargs)
        self.app.closeAllWindows()

    def showPreferences(self):
        '''
        Shows the preferences dialog.
        '''
        PreferencesDialog(self.preferences, self.__style_path_root, self.__magnitude_model.limits, parent=self).exec()

    def update_reference_series(self, names, combo, primary=True):
        '''
        Updates the reference series dropdown with the current curve names.
        '''
        current_reference = combo.currentText()
        try:
            combo.blockSignals(True)
            combo.clear()
            combo.addItem('None')
            for name in names:
                combo.addItem(name)
            idx = combo.findText(current_reference)
            if idx != -1:
                combo.setCurrentIndex(idx)
            else:
                self.__magnitude_model.normalise(primary=primary)
        finally:
            combo.blockSignals(False)

    def showAbout(self):
        msg_box = QMessageBox()
        msg_box.setText(
            f"<a href='https://github.com/3ll3d00d/beqdesigner'>BEQ Designer</a> v{self.__version} by 3ll3d00d")
        msg_box.setIcon(QMessageBox.Information)
        msg_box.setWindowTitle('About')
        msg_box.exec()

    def showExtractAudioDialog(self):
        '''
        Show the extract audio dialog.
        '''
        from model.extract import ExtractAudioDialog
        ExtractAudioDialog(self, self.preferences, self.__signal_model).show()

    def showRemuxAudioDialog(self):
        '''
        Show the remux audio dialog.
        '''
        from model.extract import ExtractAudioDialog
        ExtractAudioDialog(self, self.preferences, self.__signal_model,
                           default_signal=self.__default_signal, is_remux=True).show()

    def showBatchExtractDialog(self):
        '''
        Show the batch extract dialog.
        '''
        from model.batch import BatchExtractDialog
        BatchExtractDialog(self, self.preferences).show()

    def showAnalyseAudioDialog(self):
        '''
        Show the analyse audio dialog.
        '''
        from model.analysis import AnalyseSignalDialog
        AnalyseSignalDialog(self.preferences, self.__signal_model).exec()

    def showExportFRDDialog(self):
        '''
        Shows the export frd dialog.
        '''
        from model.export import ExportSignalDialog
        ExportSignalDialog(self.preferences, self.__signal_model, self, self.statusbar).exec()

    def showExportSignalDialog(self):
        '''
        Shows the export signal dialog.
        '''
        from model.export import ExportSignalDialog, Mode
        ExportSignalDialog(self.preferences, self.__signal_model, self, self.statusbar, mode=Mode.SIGNAL).exec()

    def exportProject(self):
        '''
        Exports the project to a file.
        '''
        file_name = QFileDialog(self).getSaveFileName(self, 'Export Project', f"project.beq", "BEQ Project (*.beq)")
        file_name = str(file_name[0]).strip()
        if len(file_name) > 0:
            output = self.__signal_model.to_json()
            if not file_name.endswith('.beq'):
                file_name += '.beq'
            with gzip.open(file_name, 'wb+') as outfile:
                outfile.write(json.dumps(output).encode('utf-8'))
            self.statusbar.showMessage(f"Saved project to {file_name}")

    def importProject(self):
        '''
        Allows the user to load a fresh project.
        '''

        def parser(file_name):
            with gzip.open(file_name, 'r') as infile:
                return json.loads(infile.read().decode('utf-8'))
        input = self.__load('*.beq', 'Load Project', parser)
        if input is not None:
            with (wait_cursor()):
                from model.codec import signalmodel_from_json
                self.__signal_model.replace(signalmodel_from_json(input, self.preferences))
                self.__magnitude_model.redraw()

    def normaliseSignalMagnitude(self):
        '''
        Handles reference series change.
        '''
        if self.signalReference.currentText() == 'None':
            self.__magnitude_model.normalise(primary=True)
        else:
            self.__magnitude_model.normalise(primary=True, curve=self.signalReference.currentText())

    def normaliseFilterMagnitude(self):
        '''
        Handles reference series change.
        '''
        if self.filterReference.currentText() == 'None':
            self.__magnitude_model.normalise(primary=False)
        else:
            self.__magnitude_model.normalise(primary=False, curve=self.filterReference.currentText())

    def showLimits(self):
        '''
        Shows the limits dialog for the main chart.
        '''
        self.__magnitude_model.show_limits()

    def showValues(self):
        '''
        Shows the values dialog for the main chart.
        '''
        self.__magnitude_model.show_values()

    def changeFilterVisibility(self, selected_filters):
        '''
        Changes which filters are visible on screen.
        '''
        self.preferences.set(DISPLAY_SHOW_FILTERS, selected_filters)
        self.__filter_model.post_update()
        self.__magnitude_model.redraw()

    def changeSignalVisibility(self, selected_signals):
        '''
        Changes which signals are visible on screen.
        '''
        self.preferences.set(DISPLAY_SHOW_SIGNALS, selected_signals)
        self.__signal_model.post_update()
        self.__magnitude_model.redraw()

    def changeSignalFilterVisibility(self, selected):
        '''
        Changes which filtered signals are visible on screen.
        '''
        self.preferences.set(DISPLAY_SHOW_FILTERED_SIGNALS, selected)
        self.__signal_model.post_update()
        self.__magnitude_model.redraw()

    def changeLegendVisibility(self):
        '''
        Changes whether the legend is visible.
        '''
        self.preferences.set(DISPLAY_SHOW_LEGEND, self.showLegend.isChecked())
        self.__magnitude_model.redraw()

    def toggleTilt(self):
        ''' applies a 3dB/octave tilt to the signals '''
        self.__signal_model.tilt(self.equalEnergyTilt.isChecked())
        self.__magnitude_model.redraw()

    def applyPreset1(self):
        '''
        Applies the preset to the model.
        '''
        self.__apply_preset(1)

    def applyPreset2(self):
        '''
        Applies the preset to the model.
        '''
        self.__apply_preset(2)

    def applyPreset3(self):
        '''
        Applies the preset to the model.
        '''
        self.__apply_preset(3)

    def __apply_preset(self, idx):
        preset_key = FILTERS_PRESET_x % idx
        preset = self.preferences.get(preset_key)
        if preset is not None:
            self.__apply_filter(preset)
            self.__filter_model.filter.preset_idx = idx
            self.__check_active_preset(idx)
            self.on_filter_selected()

    def __check_active_preset(self, preset_idx):
        pattern = re.compile("^preset[0-9]Button$")
        for attr in [at for at in dir(self) if pattern.match(at)]:
            getattr(self, attr).setIcon(QIcon())
        if preset_idx > 0:
            if hasattr(self, f"preset{preset_idx}Button"):
                getattr(self, f"preset{preset_idx}Button").setIcon(qta.icon('fa5s.check'))
            else:
                logger.warning(f"Ignoring attempt to activate an unknown preset {preset_idx}")

    def enable_preset(self, idx):
        preset_key = FILTERS_PRESET_x % idx
        getattr(self, f"preset{idx}Button").setEnabled(self.preferences.has(preset_key))

    def load_preset(self, idx):
        '''
        Allows the user to load a preset from a saved filter.
        :param idx: the index.
        '''
        from model.filter import load_filter

        def __load_preset():
            preset_key = FILTERS_PRESET_x % idx
            input = load_filter(self, self.statusbar)
            if input is not None:
                logger.info(f"Loaded filter for preset {idx}")
                self.preferences.set(preset_key, input)
            self.enable_preset(idx)

        return __load_preset

    def set_preset(self, idx):
        '''
        Saves the current filter to the numbered preset.
        :param idx: the index.
        '''

        def __set_preset():
            preset_key = FILTERS_PRESET_x % idx
            if len(self.__filter_model) > 0:
                input = self.__filter_model.filter.to_json()
                self.preferences.set(preset_key, input)
                self.enable_preset(idx)

        return __set_preset

    def clear_preset(self, idx):
        '''
        Yields a function which will clears the specified preset.
        :param idx: the preset index.
        '''
        prefs = self.preferences
        button = getattr(self, f"preset{idx}Button")

        def __clear_preset():
            prefs.set(FILTERS_PRESET_x % idx, None)
            button.setEnabled(False)

        return __clear_preset

    def add_beq_filter(self):
        '''
        Presents a file dialog to the user so they can choose a minidsp beq filter to load.
        :return: the loaded filter, if any.
        '''
        from model.minidsp import load_as_filter

        filters, filt_file = load_as_filter(self, self.preferences, self.__get_selected_signal().fs)
        if filters is not None:
            self.__load_filter(os.path.splitext(os.path.split(filt_file)[1])[0], filters)
            self.statusbar.showMessage(f"Loaded {filt_file}", 15000)

    def __load_filter(self, description, filters):
        self.clearFilters()
        self.__filter_model.filter.description = description
        for f in filters:
            self.__filter_model.save(f)

    def show_catalogue(self):
        '''
        Shows the catalogue dialog.
        '''
        from model.catalogue import CatalogueDialog
        dialog = CatalogueDialog(self, self.preferences, self.__load_filter)
        dialog.show()
        dialog.setVisible(True)

    def merge_minidsp_xml(self):
        '''
        Shows the merge minidsp XML dialog.
        '''
        from model.merge import MergeFiltersDialog
        dialog = MergeFiltersDialog(self, self.preferences, self.statusbar)
        dialog.exec()

    def create_avs_post(self):
        '''
        Shows the create post dialog.
        '''
        from model.postbuilder import CreateAVSPostDialog
        dialog = CreateAVSPostDialog(self, self.preferences, self.__filter_model, self.__get_selected_signal())
        dialog.exec()

    def sync_htp1(self):
        '''
        displays the Sync HTP1 dialog
        '''
        from model.sync import SyncHTP1Dialog
        dialog = SyncHTP1Dialog(self, self.preferences, self.__signal_model)
        dialog.exec()

    def __import_jriver_peq(self):
        from model.jriver.ui import JRiverDSPDialog
        JRiverDSPDialog(self, self.preferences).exec()

    def show_y2_tool_buttons(self, show=True):
        self.mainChartRightTools.setVisible(show)

    def y2_max_plus_10(self):
        self.__magnitude_model.limits.shift(y2_max=10)

    def y2_max_plus_5(self):
        self.__magnitude_model.limits.shift(y2_max=5)

    def y2_max_minus_5(self):
        self.__magnitude_model.limits.shift(y2_max=-5)

    def y2_max_minus_10(self):
        self.__magnitude_model.limits.shift(y2_max=-10)

    def y2_auto_on(self):
        auto = self.y2AutoOnButton.isChecked()
        self.__magnitude_model.limits.y2_auto = auto
        with block_signals(self.y2AutoOffButton):
            self.y2AutoOffButton.setChecked(not auto)

    def y2_auto_off(self):
        auto = self.y2AutoOffButton.isChecked()
        self.__magnitude_model.limits.y2_auto = not auto
        with block_signals(self.y2AutoOnButton):
            self.y2AutoOnButton.setChecked(not auto)

    def y2_min_plus_10(self):
        self.__magnitude_model.limits.shift(y2_min=10)

    def y2_min_plus_5(self):
        self.__magnitude_model.limits.shift(y2_min=5)

    def y2_min_minus_5(self):
        self.__magnitude_model.limits.shift(y2_min=-5)

    def y2_min_minus_10(self):
        self.__magnitude_model.limits.shift(y2_min=-10)

    def y1_max_plus_10(self):
        self.__magnitude_model.limits.shift(y1_max=10)

    def y1_max_plus_5(self):
        self.__magnitude_model.limits.shift(y1_max=5)

    def y1_max_minus_5(self):
        self.__magnitude_model.limits.shift(y1_max=-5)

    def y1_max_minus_10(self):
        self.__magnitude_model.limits.shift(y1_max=-10)

    def y1_auto_on(self):
        auto = self.y1AutoOnButton.isChecked()
        self.__magnitude_model.limits.y1_auto = auto
        with block_signals(self.y1AutoOffButton):
            self.y1AutoOffButton.setChecked(not auto)

    def y1_auto_off(self):
        auto = self.y1AutoOffButton.isChecked()
        self.__magnitude_model.limits.y1_auto = not auto
        with block_signals(self.y1AutoOnButton):
            self.y1AutoOnButton.setChecked(not auto)

    def y1_min_plus_10(self):
        self.__magnitude_model.limits.shift(y1_min=10)

    def y1_min_plus_5(self):
        self.__magnitude_model.limits.shift(y1_min=5)

    def y1_min_minus_5(self):
        self.__magnitude_model.limits.shift(y1_min=-5)

    def y1_min_minus_10(self):
        self.__magnitude_model.limits.shift(y1_min=-10)

    def export_beq_filter(self):
        file_name = QFileDialog(self).getSaveFileName(self, 'Export BEQ Filter', f"beq.xml", "XML (*.xml)")
        file_name = str(file_name[0]).strip()
        if len(file_name) > 0:
            if getattr(sys, 'frozen', False):
                file_path = os.path.join(sys._MEIPASS, 'flat24hd.xml')
            else:
                file_path = os.path.abspath(os.path.join(os.path.dirname('__file__'), '../xml/flat24hd.xml'))
            from model.minidsp import HDXmlParser
            from model.merge import DspType
            parser = HDXmlParser(DspType.MINIDSP_TWO_BY_FOUR_HD, False)
            output_xml, _ = parser.convert(file_path, self.__filter_model.filter)
            with open(file_name, 'w+') as f:
                f.write(output_xml)

    def __open_merge_signal_dialog(self):
        from model.signal import MergeSignalDialog
        MergeSignalDialog(self.preferences, self.__signal_model, parent=self).exec()

class SaveChartDialog(QDialog, Ui_saveChartDialog):
    '''
    Save Chart dialog
    '''

    def __init__(self, parent, name, figure, processor, statusbar=None):
        super(SaveChartDialog, self).__init__(parent)
        self.setupUi(self)
        self.name = name
        self.figure = figure
        self.processor = processor
        self.__x, self.__y = processor.get_dims(self.figure)
        self.__aspectRatio = self.__x / self.__y
        self.widthPixels.setValue(self.__x)
        self.heightPixels.setValue(self.__y)
        self.statusbar = statusbar
        self.__dialog = QFileDialog(parent=self)

    def accept(self):
        formats = "Portable Network Graphic (*.png)"
        file_name = self.__dialog.getSaveFileName(self, 'Export Chart', f"{self.name}.png", formats)
        if file_name:
            output_file = str(file_name[0]).strip()
            if len(output_file) == 0:
                return
            else:
                self.processor.export(self.figure, self.widthPixels.value(), output_file)
                if self.statusbar is not None:
                    self.statusbar.showMessage(f"Saved {self.name} to {output_file}", 5000)
        QDialog.accept(self)

    def set_height(self, newWidth):
        '''
        Updates the height as the width changes according to the aspect ratio.
        :param newWidth: the new width.
        '''
        self.heightPixels.setValue(int(math.floor(newWidth / self.__aspectRatio)))

class MatplotlibExportProcessor:

    def __init__(self, figure):
        self.__figure = figure
        self.__x, self.__y = figure.get_size_inches() * figure.dpi

    def get_dims(self, plot):
        return self.__x, self.__y

    def export(self, plot, width, output_file):
        scale_factor = width / self.__x
        plot.savefig(output_file, format='png', dpi=plot.dpi * scale_factor)

class PyQtGraphExportProcessor:

    @staticmethod
    def get_dims(plot):
        return plot.size().width(), plot.size().height()

    def export(self, plot, width, output_file):
        from pyqtgraph.exporters import ImageExporter
        exporter = ImageExporter(plot.getPlotItem())
        exporter.parameters()['width'] = width
        self.__force_to_int('height', exporter)
        self.__force_to_int('width', exporter)
        exporter.export(output_file)

    @staticmethod
    def __force_to_int(param_name, exporter):
        h = exporter.params.param(param_name)
        orig_h = int(exporter.parameters()[param_name])
        with block_signals(h):
            h.setValue(orig_h + 0.1)
            h.setValue(orig_h)

class ExportBiquadDialog(QDialog, Ui_exportBiquadDialog):
    '''
    Export Biquads Dialog
    '''

    def __init__(self, filter, prefs):
        super(ExportBiquadDialog, self).__init__()
        self.setupUi(self)
        self.__filter = filter
        self.__prefs = prefs
        self.setDefaults.setIcon(qta.icon('fa5s.save'))
        self.copyToClipboardBtn.setIcon(qta.icon('fa5s.clipboard'))
        self.saveToFile.setIcon(qta.icon('fa5s.file-export'))
        self.maxBiquads.setValue(prefs.get(BIQUAD_EXPORT_MAX))
        self.fs.setCurrentText(prefs.get(BIQUAD_EXPORT_FS))
        self.outputFormat.setCurrentText(prefs.get(BIQUAD_EXPORT_DEVICE))
        self.update_format(self.outputFormat.currentText())
        self.outputFormat.addItem('Minidsp DDRC-24')

    def update_format(self, selected_format):
        if selected_format == 'User Selected':
            self.fs.setVisible(True)
            self.fsLabel.setVisible(True)
            self.fs.setEnabled(True)
            self.maxBiquads.setVisible(True)
            self.maxBiquads.setEnabled(True)
            self.maxBiquadsLabel.setVisible(True)
            self.maxBiquads.setMaximum(1000)
            self.maxBiquads.setMinimum(1)
            self.maxBiquads.setSingleStep(1)
            self.showHex.setVisible(False)
        elif selected_format == 'Equalizer APO':
            self.fs.setVisible(False)
            self.fsLabel.setVisible(False)
            self.maxBiquads.setVisible(False)
            self.maxBiquadsLabel.setVisible(False)
            self.showHex.setVisible(False)
        else:
            if selected_format == 'Minidsp 2x4HD':
                self.fs.setCurrentText('96000')
                self.maxBiquads.setMaximum(10)
                self.maxBiquads.setMinimum(1)
                self.maxBiquads.setValue(10)
                self.maxBiquads.setSingleStep(1)
            else:
                self.fs.setCurrentText('48000')
                value = 10 if selected_format == 'Minidsp DDRC-24' else 6
                self.maxBiquads.setMaximum(value)
                self.maxBiquads.setMinimum(value)
                self.maxBiquads.setValue(value)
                self.maxBiquads.setSingleStep(2)
            self.showHex.setVisible('HD' in selected_format)
            self.fs.setEnabled(False)
            self.fs.setVisible(True)
            self.maxBiquads.setVisible(True)
            self.maxBiquadsLabel.setVisible(True)
            self.fsLabel.setVisible(True)
            self.maxBiquads.setEnabled(True)
        if not self.showHex.isVisible():
            self.showHex.setChecked(False)
        self.updateBiquads()

    def updateBiquads(self):
        has_txt = False
        if self.__filter is not None and len(self.__filter) > 0:
            if 'Equalizer APO' == self.outputFormat.currentText():
                text = "\n".join([f"Filter {idx+1}: {t}" for idx, t in enumerate(list(flatten([as_equalizer_apo(f) for f in self.__filter])))])
                if self.__filter.description != COMBINED:
                    text = f"# {self.__filter.description}\n{text}"
            else:
                self.__filter = self.__filter.resample(int(self.fs.currentText()))
                is_fixed_point = 'Minidsp 10x10HD' == self.outputFormat.currentText() or 'Minidsp 2x4' == self.outputFormat.currentText()
                kwargs = {'to_hex': self.showHex.isChecked(), 'fixed_point': is_fixed_point}
                biquads = list(flatten([self.__filter.format_biquads(self.outputFormat.currentText() != 'User Selected',
                                                                     **kwargs)]))
                if len(biquads) < self.maxBiquads.value():
                    passthrough = [Passthrough()] * (self.maxBiquads.value() - len(biquads))
                    biquads += list(flatten([x.format_biquads(self.outputFormat.currentText() != 'User Selected', **kwargs) for x in passthrough]))
                prefix = 'hex' if self.showHex.isChecked() else 'biquad'
                text = ",\n".join([f"{prefix}{idx+1},\n{bq}" for idx, bq in enumerate(biquads)])
            has_txt = len(text.strip()) > 0
            self.biquads.setPlainText(text)
        if not has_txt:
            self.copyToClipboardBtn.setEnabled(False)
            self.saveToFile.setEnabled(False)

    def save(self):
        ''' saves the current preferences '''
        self.__prefs.set(BIQUAD_EXPORT_FS, self.fs.currentText())
        self.__prefs.set(BIQUAD_EXPORT_MAX, self.maxBiquads.value())
        self.__prefs.set(BIQUAD_EXPORT_DEVICE, self.outputFormat.currentText())

    def copyToClipboard(self):
        ''' copies all the text to the clipboard '''
        self.biquads.selectAll()
        self.biquads.copy()
        cursor = self.biquads.textCursor()
        cursor.movePosition(QTextCursor.End)
        self.biquads.setTextCursor(cursor)

    def export(self):
        txt = str(self.biquads.toPlainText()).strip()
        if len(txt) > 0:
            file_name = QFileDialog(self).getSaveFileName(self, 'Export Biquads', f"biquads.txt", "Txt (*.txt)")
            file_name = str(file_name[0]).strip()
            if len(file_name) > 0:
                with open(file_name, 'w') as f:
                    f.write(txt)

def flatten(l):
    '''
    flatten an irregularly shaped list of lists (of lists of lists...)
    solution from https://stackoverflow.com/questions/2158395/flatten-an-irregular-list-of-lists
    '''
    for el in l:
        if isinstance(el, abc.Iterable) and not isinstance(el, (str, bytes)):
            yield from flatten(el)
        else:
            yield el

def make_app():
    app = QApplication(sys.argv)
    if getattr(sys, 'frozen', False):
        icon_path = os.path.join(sys._MEIPASS, 'Icon.ico')
    else:
        icon_path = os.path.abspath(os.path.join(os.path.dirname('__file__'), '../icons/Icon.ico'))
    if os.path.exists(icon_path):
        app.setWindowIcon(QIcon(icon_path))
    prefs = Preferences(QSettings("3ll3d00d", "beqdesigner"))
    if prefs.get(STYLE_MATPLOTLIB_THEME) == f"{STYLE_MATPLOTLIB_THEME_DEFAULT}_extra":
        import qdarkstyle
        style = qdarkstyle.load_stylesheet_from_environment()
        app.setStyleSheet(style)
    return app, prefs

if __name__ == '__main__':
    app, prefs = make_app()
    form = BeqDesigner(app, prefs)
    # setup the error handler
    e_dialog = QErrorMessage(form)
    e_dialog.setWindowModality(QtCore.Qt.WindowModal)
    font = QFont()
    font.setFamily("Consolas")
    font.setPointSize(8)
    e_dialog.setFont(font)
    # add the exception handler so we can see the errors in a QErrorMessage
    sys._excepthook = sys.excepthook

    def dump_exception_to_log(exctype, value, tb):
        import traceback
        print(exctype, value, tb)
        global e_dialog
        if e_dialog is not None:
            formatted = traceback.format_exception(etype=exctype, value=value, tb=tb)
            e_dialog.setWindowTitle('Unexpected Error')
            url = 'https://github.com/3ll3d00d/beqdesigner/issues/new'
            msg = f"Unexpected Error detected, go to {url} to log the issue<p>{'<br>'.join(formatted)}"
            e_dialog.showMessage(msg)
            e_dialog.resize(1200, 400)

    sys.excepthook = dump_exception_to_log

    # show the form and exec the app
    form.show()
    app.exec_()

class PlotWidgetWithDateAxis(pg.PlotWidget):
    def __init__(self, parent=None, background='default', **kargs):
        super().__init__(parent=parent,
                         background=background,
                         axisItems={'bottom': TimeAxisItem(orientation='bottom')},
                         **kargs)

class TimeAxisItem(pg.AxisItem):
    def tickStrings(self, values, scale, spacing):
        import datetime
        return [str(datetime.timedelta(seconds=value)).split('.')[0] for value in values]

class NoFillDoubleSpinBox(QtWidgets.QDoubleSpinBox):

    def __init__(self, parent=None):
        super().__init__(parent)

    def textFromValue(self, p_float: float):
        txt = super().textFromValue(p_float).rstrip('0')
        return f"{txt}0" if p_float.is_integer() else txt


from __future__ import division
import os.path
from .listdataset import  ListDataset

import numpy as np
import random
import flow_transforms
import pdb
import glob
from tqdm import tqdm
from tqdm import trange

try:
    import cv2
except ImportError as e:
    import warnings
    with warnings.catch_warnings():
        warnings.filterwarnings("default", category=ImportWarning)
        warnings.warn("failed to load openCV, which is needed"
                      "for KITTI which uses 16bit PNG images", ImportWarning)

def index_items(sample_list, train_index, val_index):
    train_samples = [sample_list[idx] for idx in train_index]
    val_samples = [sample_list[idx] for idx in val_index]

    return train_samples, val_samples

def make_dataset(path):
    # we train and val seperately to tune the hyper-param and use all the data for the final training
    #train_list_path = os.path.join(dir, 'train.txt') # use train_Val.txt for final report
    #val_list_path = os.path.join(dir, 'val.txt')
    gt_nii_img = glob.glob(path + '*GG/seg/*nii.gz')
    t1_nii_img = [item.replace('/seg/','/t1/').replace('_seg.', '_t1.') for item in gt_nii_img]
    t1ce_nii_img = [item.replace('/seg/','/t1ce/').replace('_seg.', '_t1ce.') for item in gt_nii_img]
    t2_nii_img = [item.replace('/seg/','/t2/').replace('_seg.', '_t2.') for item in gt_nii_img]
    flair_nii_img = [item.replace('/seg/','/flair/').replace('_seg.', '_flair.') for item in gt_nii_img]

    total_index = list(range(len(gt_nii_img)))
    #total_index = list(range(20))

    random.shuffle(total_index)
    train_num = int(len(total_index) * 0.75*0.85)
    val_num = int(len(total_index) * 0.75*0.15)
    train_index = total_index[:train_num]
    val_index = total_index[train_num:(train_num+val_num)]

    train_gt_list, val_gt_list = index_items(gt_nii_img, train_index, val_index)
    train_t1_list, val_t1_list = index_items(t1_nii_img, train_index, val_index)
    train_t1ce_list, val_t1ce_list = index_items(t1ce_nii_img, train_index, val_index)
    train_t2_list, val_t2_list = index_items(t2_nii_img, train_index, val_index)
    train_flair_list, val_flair_list = index_items(flair_nii_img, train_index, val_index)

    train_dict = {
        't1':train_t1_list,
        't1ce':train_t1ce_list,
        't2':train_t2_list,
        'flair':train_flair_list,
        'gt':train_gt_list
        }

    val_dict = {
        't1':val_t1_list,
        't1ce':val_t1ce_list,
        't2':val_t2_list,
        'flair':val_flair_list,
        'gt':val_gt_list
        }
    return train_dict, val_dict

def BSD_loader(data_dict):
    # cv2.imread is faster than io.imread usually
    nonzero_index = []
    def pick_slice(seg_im):
        h,w,slice_num = seg_im.shape
        seg_im = np.reshape(seg_im, (-1, slice_num))
        sum_value  = np.sum(seg_im, axis=0)
        nonzero_index = np.where(sum_value > 0)

        return nonzero_index[0]

    seg_list = data_dict['gt']
    t1_list = data_dict['t1']
    t1ce_list = data_dict['t1ce']
    t2_list = data_dict['t2']
    flair_list = data_dict['flair']

    nii_num = len(seg_list)
    im_list = []
    gt_list = []
    for i in trange(nii_num):
        seg_im = nib.load(seg_list[i]).get_fdata()
        nonzero_index = pick_slice(seg_im)
        seg_slice = seg_im[:,:,nonzero_index]
        seg_slice = np.transpose(seg_slice, (2, 0, 1)).astype(np.uint8)

        t1_im = nib.load(t1_list[i]).get_fdata()
        t1_slice = t1_im[:,:,nonzero_index]
        t1_slice = np.transpose(t1_slice, (2, 0, 1))
        #t1_slice = (t1_slice / np.max(t1_slice) * 255).astype(np.uint8)

        t1ce_im = nib.load(t1ce_list[i]).get_fdata()
        t1ce_slice = t1ce_im[:,:,nonzero_index]
        t1ce_slice = np.transpose(t1ce_slice, (2, 0, 1))
        #t1ce_slice = (t1ce_slice / np.max(t1ce_slice) * 255).astype(np.uint8)

        t2_im = nib.load(t2_list[i]).get_fdata()
        t2_slice = t2_im[:,:,nonzero_index]
        t2_slice = np.transpose(t2_slice, (2, 0, 1))
        #t2_slice = (t2_slice / np.max(t2_slice) * 255).astype(np.uint8)

        flair_im = nib.load(flair_list[i]).get_fdata()
        flair_slice = flair_im[:,:,nonzero_index]
        flair_slice = np.transpose(flair_slice, (2, 0, 1))
        #flair_slice = (flair_slice / np.max(flair_slice) * 255).astype(np.uint8)

        img_data = np.concatenate([t1_slice[:,:,:,np.newaxis], t1ce_slice[:,:,:,np.newaxis], t2_slice[:,:,:,np.newaxis], flair_slice[:,:,:,np.newaxis]], axis=-1)
        img_data = (img_data / np.max(img_data)*255).astype(np.uint8)

        im_list.append(img_data)
        gt_list.append(seg_slice[:,:,:,np.newaxis])

    im_data = np.concatenate(im_list, axis=0)
    gt_data = np.concatenate(gt_list, axis=0)

    brain = (im_data[:,:,:,:1] > 0) * 3
    mask = 1 - (gt_data > 0) * 1
    normal_brain = brain * mask
    gt_data = normal_brain.astype(np.uint8) + gt_data.astype(np.uint8)

    return im_data, gt_data
    #return np.transpose(im_data, (0,3,1,2)), np.transpose(gt_data, (0,3,1,2))

def BSD500(root, transform=None, target_transform=None, val_transform=None,
              co_transform=None, split=None):
    train_dict, val_dict= make_dataset(root)

    if val_transform ==None:
        val_transform = transform

    train_dataset = ListDataset(root, 'bsd500', train_dict, transform,
                                target_transform, co_transform,
                                loader=BSD_loader, datatype = 'train')

    val_dataset = ListDataset(root, 'bsd500', val_dict, val_transform,
                               target_transform, flow_transforms.CenterCrop((240,240)),
                               loader=BSD_loader, datatype = 'val')

    return train_dataset, val_dataset

def mix_datasets(args):
    train_dataset = ListDataset(args, 'train')
    val_dataset = ListDataset(args, 'val')

    return train_dataset, val_dataset


import argparse
import math
import random
import os

import numpy as np
import torch
from torch import nn, autograd, optim
from torch.nn import functional as F
from torch.utils import data
import torch.distributed as dist
from torchvision import transforms, utils
from tqdm import tqdm
from tensorfn import load_arg_config, distributed as dist, get_logger

try:
    import wandb

except ImportError:
    wandb = None

from stylegan2.dataset import MultiResolutionDataset
from stylegan2.distributed import (
    get_rank,
    synchronize,
    reduce_loss_dict,
    reduce_sum,
    get_world_size,
)
from stylegan2.op import conv2d_gradfix
from stylegan2.non_leaking import augment, AdaptiveAugment
from stylegan2.model import Discriminator
from config import GANConfig
from model import Generator, filter_parameters

def data_sampler(dataset, shuffle, distributed):
    if distributed:
        return data.distributed.DistributedSampler(dataset, shuffle=shuffle)

    if shuffle:
        return data.RandomSampler(dataset)

    else:
        return data.SequentialSampler(dataset)

def requires_grad(model, flag=True):
    for p in model.parameters():
        p.requires_grad = flag

@torch.no_grad()
def accumulate(model1, model2, decay=0.999):
    par1 = dict(model1.named_parameters())
    par2 = dict(model2.named_parameters())

    for k in par1.keys():
        par1[k].data.mul_(decay).add_(par2[k].data, alpha=1 - decay)

    buf1 = dict(model1.named_buffers())
    buf2 = dict(model2.named_buffers())

    for k in buf1.keys():
        if "ema_var" not in k:
            continue

        buf1[k].data.mul_(0).add_(buf2[k].data, alpha=1)

def sample_data(loader):
    while True:
        for batch in loader:
            yield batch

def d_logistic_loss(real_pred, fake_pred):
    real_loss = F.softplus(-real_pred)
    fake_loss = F.softplus(fake_pred)

    return real_loss.mean() + fake_loss.mean()

def d_r1_loss(real_pred, real_img):
    with conv2d_gradfix.no_weight_gradients():
        grad_real, = autograd.grad(
            outputs=real_pred.sum(), inputs=real_img, create_graph=True
        )
    grad_penalty = grad_real.pow(2).reshape(grad_real.shape[0], -1).sum(1).mean()

    return grad_penalty

def g_nonsaturating_loss(fake_pred):
    loss = F.softplus(-fake_pred).mean()

    return loss

def g_path_regularize(fake_img, latents, mean_path_length, decay=0.01):
    noise = torch.randn_like(fake_img) / math.sqrt(
        fake_img.shape[2] * fake_img.shape[3]
    )
    grad, = autograd.grad(
        outputs=(fake_img * noise).sum(), inputs=latents, create_graph=True
    )
    path_lengths = torch.sqrt(grad.pow(2).sum(2).mean(1))

    path_mean = mean_path_length + decay * (path_lengths.mean() - mean_path_length)

    path_penalty = (path_lengths - path_mean).pow(2).mean()

    return path_penalty, path_mean.detach(), path_lengths

def make_noise(batch, latent_dim, n_noise, device):
    if n_noise == 1:
        return torch.randn(batch, latent_dim, device=device)

    noises = torch.randn(n_noise, batch, latent_dim, device=device).unbind(0)

    return noises

def mixing_noise(batch, latent_dim, prob, device):
    if prob > 0 and random.random() < prob:
        return make_noise(batch, latent_dim, 2, device)

    else:
        return [make_noise(batch, latent_dim, 1, device)]

def set_grad_none(model, targets):
    for n, p in model.named_parameters():
        if n in targets:
            p.grad = None

def train(conf, loader, generator, discriminator, g_optim, d_optim, g_ema, device):
    loader = sample_data(loader)

    pbar = range(conf.training.iter)

    if get_rank() == 0:
        pbar = tqdm(
            pbar, initial=conf.training.start_iter, dynamic_ncols=True, smoothing=0.01
        )

    mean_path_length = 0

    d_loss_val = 0
    r1_loss = torch.tensor(0.0, device=device)
    g_loss_val = 0
    path_loss = torch.tensor(0.0, device=device)
    path_lengths = torch.tensor(0.0, device=device)
    mean_path_length_avg = 0
    loss_dict = {}

    if conf.distributed:
        g_module = generator.module
        d_module = discriminator.module

    else:
        g_module = generator
        d_module = discriminator

    accum = 0.5 ** (32 / (10 * 1000))
    ada_aug_p = conf.training.augment_p if conf.training.augment_p > 0 else 0.0
    r_t_stat = 0

    if conf.training.augment and conf.training.augment_p == 0:
        ada_augment = AdaptiveAugment(
            conf.training.ada_target, conf.training.ada_length, 8, device
        )

    sample_z = torch.randn(
        conf.training.n_sample, conf.generator["style_dim"], device=device
    )

    for idx in pbar:
        i = idx + conf.training.start_iter

        if i > conf.training.iter:
            print("Done!")

            break

        real_img = next(loader)
        real_img = real_img.to(device)

        requires_grad(generator, False)
        requires_grad(discriminator, True)

        noise = make_noise(conf.training.batch, conf.generator["style_dim"], 1, device)
        generator.eval()
        fake_img = generator(noise)
        generator.train()

        if conf.training.augment:
            real_img_aug, _ = augment(real_img, ada_aug_p)
            fake_img, _ = augment(fake_img, ada_aug_p)

        else:
            real_img_aug = real_img

        fake_pred = discriminator(fake_img)
        real_pred = discriminator(real_img_aug)
        d_loss = d_logistic_loss(real_pred, fake_pred)

        loss_dict["d"] = d_loss
        loss_dict["real_score"] = real_pred.mean()
        loss_dict["fake_score"] = fake_pred.mean()

        discriminator.zero_grad()
        d_loss.backward()
        d_optim.step()

        if conf.training.augment and conf.training.augment_p == 0:
            ada_aug_p = ada_augment.tune(real_pred)
            r_t_stat = ada_augment.r_t_stat

        d_regularize = i % conf.training.d_reg_every == 0

        if d_regularize:
            real_img.requires_grad = True

            if conf.training.augment:
                real_img_aug, _ = augment(real_img, ada_aug_p)

            else:
                real_img_aug = real_img

            real_pred = discriminator(real_img_aug)
            r1_loss = d_r1_loss(real_pred, real_img)

            discriminator.zero_grad()
            (
                conf.training.r1 / 2 * r1_loss * conf.training.d_reg_every
                + 0 * real_pred[0]
            ).backward()

            d_optim.step()

        loss_dict["r1"] = r1_loss

        requires_grad(generator, True)
        requires_grad(discriminator, False)

        noise = make_noise(conf.training.batch, conf.generator["style_dim"], 1, device)
        fake_img = generator(noise)

        if conf.training.augment:
            fake_img, _ = augment(fake_img, ada_aug_p)

        fake_pred = discriminator(fake_img)
        g_loss = g_nonsaturating_loss(fake_pred)

        loss_dict["g"] = g_loss

        generator.zero_grad()
        g_loss.backward()
        g_optim.step()

        accumulate(g_ema, g_module, accum)

        loss_reduced = reduce_loss_dict(loss_dict)

        d_loss_val = loss_reduced["d"].mean().item()
        g_loss_val = loss_reduced["g"].mean().item()
        r1_val = loss_reduced["r1"].mean().item()
        real_score_val = loss_reduced["real_score"].mean().item()
        fake_score_val = loss_reduced["fake_score"].mean().item()

        if get_rank() == 0:
            pbar.set_description(
                (
                    f"d: {d_loss_val:.4f}; g: {g_loss_val:.4f}; r1: {r1_val:.4f}; "
                    f"augment: {ada_aug_p:.4f}"
                )
            )

            if wandb and conf.wandb:
                wandb.log(
                    {
                        "Generator": g_loss_val,
                        "Discriminator": d_loss_val,
                        "Augment": ada_aug_p,
                        "Rt": r_t_stat,
                        "R1": r1_val,
                        "Real Score": real_score_val,
                        "Fake Score": fake_score_val,
                    }
                )

            if i % 100 == 0:
                generator.zero_grad()
                discriminator.zero_grad()
                with torch.no_grad():
                    g_ema.eval()
                    sample = g_ema(sample_z).cpu()
                    utils.save_image(
                        sample,
                        f"sample/{str(i).zfill(6)}.png",
                        nrow=int(conf.training.n_sample ** 0.5),
                        normalize=True,
                        value_range=(-1, 1),
                    )
                    sample = None # cleanup memory

            if i % 10000 == 0:
                torch.save(
                    {
                        "g": g_module.state_dict(),
                        "d": d_module.state_dict(),
                        "g_ema": g_ema.state_dict(),
                        "g_optim": g_optim.state_dict(),
                        "d_optim": d_optim.state_dict(),
                        "conf": conf.dict(),
                        "ada_aug_p": ada_aug_p,
                    },
                    f"checkpoint/{str(i).zfill(6)}.pt",
                )

def main(conf):
    device = "cuda"
    conf.distributed = conf.n_gpu > 1

    logger = get_logger(mode=conf.logger)
    logger.info(conf.dict())

    generator = conf.generator.make().to(device)
    g_ema = conf.generator.make().to(device)
    discriminator = conf.discriminator.make().to(device)
    accumulate(g_ema, generator, 0)

    logger.info(generator)
    logger.info(discriminator)

    d_reg_ratio = conf.training.d_reg_every / (conf.training.d_reg_every + 1)

    g_optim = optim.Adam(generator.parameters(), lr=conf.training.lr_g, betas=(0, 0.99))
    d_optim = optim.Adam(
        discriminator.parameters(),
        lr=conf.training.lr_d * d_reg_ratio,
        betas=(0 ** d_reg_ratio, 0.99 ** d_reg_ratio),
    )

    if conf.ckpt is not None:
        logger.info(f"load model: {conf.ckpt}")

        ckpt = torch.load(conf.ckpt, map_location=lambda storage, loc: storage)

        try:
            ckpt_name = os.path.basename(conf.ckpt)
            conf.training.start_iter = int(os.path.splitext(ckpt_name)[0])

        except ValueError:
            pass

        generator.load_state_dict(ckpt["g"])
        discriminator.load_state_dict(ckpt["d"])
        g_ema.load_state_dict(ckpt["g_ema"])

        g_optim.load_state_dict(ckpt["g_optim"])
        d_optim.load_state_dict(ckpt["d_optim"])

    if conf.distributed:
        generator = nn.parallel.DistributedDataParallel(
            generator,
            device_ids=[dist.get_local_rank()],
            output_device=dist.get_local_rank(),
            broadcast_buffers=True,
        )

        discriminator = nn.parallel.DistributedDataParallel(
            discriminator,
            device_ids=[dist.get_local_rank()],
            output_device=dist.get_local_rank(),
            broadcast_buffers=False,
        )

    transform = transforms.Compose(
        [
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),
        ]
    )

    dataset = MultiResolutionDataset(conf.path, transform, conf.training.size)
    loader = data.DataLoader(
        dataset,
        batch_size=conf.training.batch,
        sampler=data_sampler(dataset, shuffle=True, distributed=conf.distributed),
        drop_last=True,
    )

    if get_rank() == 0 and wandb is not None and conf.wandb:
        wandb.init(project="alias free gan")

    train(conf, loader, generator, discriminator, g_optim, d_optim, g_ema, device)

if __name__ == "__main__":
    conf = load_arg_config(GANConfig)

    dist.launch(
        main, conf.n_gpu, conf.n_machine, conf.machine_rank, conf.dist_url, args=(conf,)
    )


import torch
from tqdm import tqdm
import numpy as np

from transformers import (
    CTRLLMHeadModel,
    CTRLTokenizer,
    GPT2LMHeadModel,
    GPT2Tokenizer,
    OpenAIGPTLMHeadModel,
    OpenAIGPTTokenizer,
    TransfoXLLMHeadModel,
    TransfoXLTokenizer,
    XLMTokenizer,
    XLMWithLMHeadModel,
    XLNetLMHeadModel,
    XLNetTokenizer,
)
MODEL_CLASSES = {
    "gpt2": (GPT2LMHeadModel, GPT2Tokenizer),
    "ctrl": (CTRLLMHeadModel, CTRLTokenizer),
    "openai-gpt": (OpenAIGPTLMHeadModel, OpenAIGPTTokenizer),
    "xlnet": (XLNetLMHeadModel, XLNetTokenizer),
    "transfo-xl": (TransfoXLLMHeadModel, TransfoXLTokenizer),
    "xlm": (XLMWithLMHeadModel, XLMTokenizer),
}

class Args:
    def __init__(self):
        self.model_type = 'gpt2'
        self.model_name_or_path = 'sberbank-ai/rugpt2large'

        self.prompt = ''
        self.length = 50
        self.stop_token = '</s>'

        self.k = 5
        self.p = .95
        self.temperature = 1

        self.repetition_penalty = 1
        self.num_return_sequences = 1

        self.device = 'cuda:0'
        self.seed = 42

def get_gpt2_ppl_corpus(test_sentences):
    args = Args()
    args.model_name_or_path = 'sberbank-ai/rugpt2large'
    model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
    tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path)
    model = model_class.from_pretrained(args.model_name_or_path)
    model.to(args.device)

    lls = []
    weights = []
    for text in tqdm(test_sentences):
        encodings = tokenizer(f'\n{text}\n', return_tensors='pt')
        input_ids = encodings.input_ids.to(args.device)
        target_ids = input_ids.clone()

        w = max(0, len(input_ids[0]) - 1)
        if w > 0:
            with torch.no_grad():
                outputs = model(input_ids, labels=target_ids)
                log_likelihood = outputs[0]
                ll = log_likelihood.item()
        else:
            ll = 0
        lls.append(ll)
        weights.append(w)

    return np.dot(lls, weights) / sum(weights)


# -*- coding: utf-8 -*-
"""
TODO:
add method for inserting nonlinear triplets, so they can be indexed for 2nd order noise analysis

have measurement objects inject objects into the solver, so that the solver state is totally separate

have coherent measurement systems able to run the solver independently to generate system-altering measurements
"""
from __future__ import division, print_function, unicode_literals
from phasor.utilities.print import print
from collections import defaultdict
from declarative import bunch

import numpy as np
import warnings

import declarative as declarative

from ..utilities.future_from_2 import super
from ..base import (
    DictKey,
    Frequency,
    SystemElementBase,
    ClassicalFreqKey,
    PTREE_ASSIGN,
)

from .matrix_algorithm import (
    MatrixBuildAlgorithm,
)

from . import ports_algorithm
from . import solver_algorithm

from ..base import Element, RootElement
from ..base.ports import PostBondKey
from ..base import ports
from ..base import visitors as VISIT

from ..optics import (
    OpticalFrequency,
    OpticalFreqKey,
)

class SystemSymbols(Element):
    """
    Separate Class for constants to allow replacing them with symbols or more math appropriate ones
    """

    def __build__(self):
        PTREE_ASSIGN(self).c_m_s                 = 299792458
        PTREE_ASSIGN(self).kB_J_K                = 1.380658e-23
        PTREE_ASSIGN(self).h_Js                  = 6.6260700408e-34
        PTREE_ASSIGN(self).hbar_Js               = 1.0545718001e-34
        PTREE_ASSIGN(self).pi                    = np.pi
        PTREE_ASSIGN(self).i                     = 1j
        PTREE_ASSIGN(self).i2pi                  = np.pi * 2j
        PTREE_ASSIGN(self).math                  = np
        PTREE_ASSIGN(self).temp_K                = 299
        super().__build__()

    def number(self, num):
        return num

class BGSystem(RootElement):
    exact_order = None
    max_N = 100
    warning_N = 20
    solver_Q_conditioning = True

    _frozen = False

    solver_name_default = 'loop_LUQ'
    #solver_name_default = 'scisparse'
    @declarative.dproperty
    def solver_name(self, val = None):
        if val is None:
            val = self.solver_name_default
        return val

    solver_symbolic_name_default = 'loop_LUQ'
    @declarative.dproperty
    def solver_symbolic_name(self, val = None):
        if val is None:
            val = self.solver_symbolic_name_default
        return val

    @declarative.dproperty
    def symbolic(self, val = False):
        return val

    @declarative.dproperty
    def include_johnson_noise(self, val = True):
        val = self.ctree.setdefault('include_johnson_noise', val)
        return val

    @declarative.dproperty
    def symbols(self):
        return SystemSymbols()

    @declarative.mproperty
    def field_space(self, FS = None):
        return FS

    @declarative.mproperty
    def freq_order_max_default(self, val = 2):
        return val

    @declarative.dproperty
    def adjust_PSD(self):
        if self.sided_spectra == 1:
            val = 2
        elif self.sided_spectra == 2:
            val = 1
        val = self.ctree.setdefault('adjust_PSD', val)
        return val

    @declarative.dproperty
    def sided_spectra(self, val = 1):
        val = self.ctree.setdefault('sided_spectra', val)
        return val

    @declarative.dproperty
    def system(self):
        #should use environment_query instead
        return self

    @declarative.mproperty
    def elements_named(self):
        return dict()

    @declarative.mproperty
    def elements(self):
        return set()

    @declarative.mproperty
    def port_owners(self):
        return dict()

    @declarative.mproperty
    def owners_ports(self):
        return dict()

    @declarative.mproperty
    def owners_ports_virtual(self):
        return dict()

    @declarative.mproperty
    def port_autoterminate(self):
        return dict()

    @declarative.mproperty
    def bond_pairs(self):
        return defaultdict(dict)

    @declarative.mproperty
    def ports_post(self):
        return dict()

    def ports_post_get(self, port):
        val = port.get(PostBondKey, 0)
        nport = port | DictKey({PostBondKey : val + 1})
        self.ports_post[port] = nport
        return nport

    @declarative.mproperty
    def ports_pre(self):
        return dict()

    def ports_pre_get(self, port):
        val = port.get(PostBondKey, 0)
        nport = port | DictKey({PostBondKey : val - 1})
        self.ports_pre[port] = nport
        return nport

    @declarative.mproperty
    def bonded_set(self):
        return set()

    @declarative.mproperty
    def port_owners_virtual(self):
        return defaultdict(set)

    @declarative.dproperty
    def F_AC(self, val = 1e-4):
        #TODO make this detect a PropertyTransforming or injected frequency object
        self.environment.own.F_AC = Frequency(
            F_Hz  = val,
            name  = 'AC',
            order = 2,
            F_center_Hz = 1e3,
        )
        return self.environment.own.F_AC

    @declarative.dproperty
    def F_carrier_1064(self):
        self.environment.own.F_carrier_1064 = OpticalFrequency(
            wavelen_m = 1064e-9,
            name = u'λIR',
        )
        return self.environment.own.F_carrier_1064

    @declarative.dproperty
    def FD_carrier_1064(self):
        return ports.FrequencyKey({self.F_carrier_1064 : 1})

    @declarative.dproperty
    def environment(self):
        return SystemElementBase()

    def optical_frequency_extract(self, key):
        iwavelen_m = 0
        freq_Hz = 0
        for F, n in list(key[ClassicalFreqKey].F_dict.items()):
            freq_Hz += n * F.F_Hz.val
        for F, n in list(key[OpticalFreqKey].F_dict.items()):
            iwavelen_m += n * F.iwavelen_m
        return iwavelen_m, freq_Hz

    @declarative.mproperty
    def fully_resolved_name_tuple(self):
        return ()

    def reject_classical_frequency_order(self, fkey):
        #TODO put this logic in a subobject
        group_N = defaultdict(lambda: 0)
        for F, N in list(fkey.F_dict.items()):
            N_limit = F.order
            if N_limit is None:
                N_limit = self.freq_order_max_default
            Na = abs(N)
            if Na > N_limit:
                return True
            #TODO revisit max frequency groups
            #for group in F.groups:
            #    group_N[group] += Na

        for group, N in list(group_N.items()):
            Nmax = self.groups_max_N.get(group, None)
            if Nmax is not None and N > Nmax:
                return True
        return False

    def reject_optical_frequency_order(self, fkey):
        #TODO put this logic in a subobject
        group_N = defaultdict(lambda: 0)
        if len(fkey.F_dict) == 0:
            #if it is truely DC then reject
            #TODO optical rectification
            return True
        for F, N in list(fkey.F_dict.items()):
            N_limit = F.order
            if N_limit is None:
                N_limit = 2
            if N > N_limit:
                return True
            if N < 1:
                return True
            #TODO revisit max frequency groups
            #for group in F.groups:
            #    group_N[group] += Na

        for group, N in list(group_N.items()):
            Nmax = self.groups_max_N.get(group, None)
            if Nmax is not None and N > Nmax:
                return True
        return False

    def classical_frequency_extract(self, key):
        #TODO put this logic in a subobject
        freq_Hz = 0
        for F, n in list(key[ClassicalFreqKey].F_dict.items()):
            freq_Hz += n * F.F_Hz.val
        return freq_Hz

    def classical_frequency_extract_center(self, ekey):
        #TODO put this logic in a subobject
        freq_Hz = 0
        for F, n in list(ekey.F_dict.items()):
            if len(np.array(F.F_center_Hz).shape) > 0:
                print("Weird F_center_Hz shape: ", F, n, F.F_center_Hz)
            freq_Hz += n * F.F_center_Hz
        return freq_Hz

    def classical_frequency_test_max(self, key, max_freq):
        if max_freq is None:
            return False
        #TODO put this logic in a subobject
        freq_Hz = 0
        for F, n in list(key[ClassicalFreqKey].F_dict.items()):
            freq_Hz += n * F.F_Hz
        if np.any(freq_Hz < max_freq):
            return False
        return True

    def _setup_sequence(self):
        self._complete()
        assert(not self._frozen)

        with self.building:
            #must do delegate calls before autotermination
            has_delegated = set()
            while True:
                no_delegates = True
                for element, delegate_call in self.targets_recurse(VISIT.bond_delegate):
                    if element not in has_delegated:
                        has_delegated.add(element)
                        delegate_call()
                        no_delegates = False
                    #print("Bond delegate", element)
                if no_delegates:
                    break

        while self._include_lst:
            self.do_includes()
            #print("INCLUDE LST: ", self._include_lst)
            self._autoterminate()

        for element in self.targets_recurse(VISIT.bond_completion):
            #print("Bond Completion", element)
            pass

        bad_completions = dict()
        while self._bond_completions_raw:
            (pe_1, pe_2), raw_port1 = self._bond_completions_raw.popitem()
            raw_port2 = self._bond_completions_raw.pop((pe_2, pe_1), None)
            if raw_port2 is None:
                bad_completions[(pe_1, pe_2)] = raw_port1
            else:
                self.bond_old(raw_port1, raw_port2)
        if bad_completions:
            raise RuntimeError("Incomplete port pairings: {0}".format(bad_completions))

        assert(not self._frozen)
        self._frozen = True

        self.port_algo = ports_algorithm.PortUpdatesAlgorithm(
            system           = self,
        )
        self.matrix_algorithm = MatrixBuildAlgorithm(
            system           = self,
            ports_algorithm  = self.port_algo,
        )
        self.solver = solver_algorithm.SystemSolver(
            system           = self,
            ports_algorithm  = self.port_algo,
            matrix_algorithm = self.matrix_algorithm,
        )
        return

    def bond_completion_raw_pair(self, pe_1, pe_2, raw_port1, raw_port2):
        self.bond_completion_raw(pe_1, pe_2, raw_port1)
        self.bond_completion_raw(pe_2, pe_1, raw_port2)

    @declarative.mproperty
    def _bond_completions_raw(self):
        return dict()

    def bond_completion_raw(self, pe_1, pe_2, raw_port):
        self._bond_completions_raw[(pe_1, pe_2)] = raw_port
        return

    def _solve(self):
        self._setup_sequence()
        self.solver.solve()
        return self.solver

    @declarative.mproperty
    def solution(self):
        return self._solve()

    @declarative.mproperty
    def _include_lst(self):
        return []

    def include(self, element):
        #TODO get rid of this deferred include business
        self._include_lst.append(element)
        #self._include(element)

    def do_includes(self):
        while self._include_lst:
            el = self._include_lst.pop()
            self._include(el)

    def port_add(self, owner, pkey):
        #TODO phase out
        self.port_owners[pkey] = owner
        self.owners_ports.setdefault(owner, []).append(pkey)
        self.owners_ports_virtual.setdefault(owner, [])
        return

    def _include(self, element):
        if self._frozen:
            raise RuntimeError("Cannot include elements or bond after solution requested")
        if element in self.elements:
            return
        #print("INCLUDE: ", element, element.name_system)
        self.elements.add(element)
        try:
            op = element.owned_ports
        except AttributeError:
            pass
        else:
            for port, pobj in list(op.items()):
                self.port_add(element, port)
        if element.name_system in self.elements_named:
            warnings.warn((
                "Multiple elements added with name {0}"
                ", may be confusing and prevent system rebuilding from functioning"
            ).format(element.name_system))
            #break the elements named so that there isn't confusion later
            self.elements_named[element.name_system] = None
        else:
            self.elements_named[element.name_system] = element

    def _autoterminate(self):
        #TODO purge this code since now bonds handle autotermination
        terminated_ports = set(self.bonded_set)
        registered_ports = set(self.port_owners.keys())
        unterminated_ports = registered_ports - terminated_ports
        if unterminated_ports and not hasattr(self, 'autoterminate'):
            self.own.autoterminate = SystemElementBase()
        for port in unterminated_ports:
            #print("UNTERMINATED: ", port)
            aterm = self.port_autoterminate.get(port, None)
            if aterm is None:
                #print("unterminated port", port)
                pass
            else:
                pobj, tclass = aterm
                #have to build within include to get the sled connection
                #TODO: make this a sane naming scheme
                #TODO GOTTA GOTTA FIX
                name = '{0}({2}-<{1}>)'.format(tclass.__name__, port, pobj.element.name_system).replace('.', '_')
                tinst = self.autoterminate.insert(tclass(), name)
                #print("Autoterminating port:", port)
                #print("with ", tinst)
                self.bond(pobj, tinst.po_Fr)

        with self.building:
            for port, obj in self.targets_recurse(VISIT.auto_terminate):
                #print("Autoterminating port:", port)
                #print("with ", obj)
                pass
        return

    def bond(self, port_A, port_B):
        #TODO ptypes testing is only for the transition to local bond scheme
        return port_A.bond(port_B)

    def bond_old(self, port_A, port_B):
        #TODO, phase this out
        self.include(port_A.element)
        self.include(port_B.element)

        either_included = False
        try:
            port_A.o
            port_B.i
        except AttributeError:
            pass
        else:
            either_included = True
            if (port_B.i in self.bonded_set):
                assert(port_B.multiple_attach)
            if (port_A.o in self.bonded_set):
                assert(port_A.multiple_attach)
            self.bonded_set.add(port_A.o)
            self.bonded_set.add(port_B.i)
            self.bond_pairs[port_A.o][port_B.i] = 1

        #also attach the duals
        try:
            port_B.o
            port_A.i
        except AttributeError:
            pass
        else:
            either_included = True
            if (port_A.i in self.bonded_set):
                assert(port_A.multiple_attach)
            if (port_B.o in self.bonded_set):
                assert(port_B.multiple_attach)
            self.bonded_set.add(port_B.o)
            self.bonded_set.add(port_A.i)
            self.bond_pairs[port_B.o][port_A.i] = 1

        if not either_included:
            raise RuntimeError("Trying to bond, but the pair is not in-out compatible")
        return

    def bond_sequence(self, first, *args):
        next = first
        for bond_arg in args[:-1]:
            self.bond(next, bond_arg)
            next = bond_arg.chain_next
            if next is None:
                raise RuntimeError("Bond in sequence doesn't indicate a natural 'next' element: ".format(bond_arg))
        self.bond(next, args[-1])
        pass

    def own_port_virtual(self, element, port):
        #TODO rename bond port virtual
        #assert(self.root is port.element.root)
        #assert(self.element.root is self.root)
        self.port_owners_virtual[port].add(element)
        self.owners_ports_virtual.setdefault(element, []).append(port)
        return

#add an additional name for the system
LinearSystem = BGSystem