#! /usr/bin/env python
#-*- coding: utf-8 -*-

"""
Documents
"""

def s(*args, **kwargs):
    import pylab
    fig = pylab.figure()
    bbox = fig.add_subplot(*args, **kwargs).get_position()
    subplot = (bbox.x0, bbox.y0, bbox.width, bbox.height)
    pylab.close(fig)
    return subplot

def draw_ps_raw(raw, figure=None, subplot=111, title='', grid=True,
                      font_family=None, tick_labels_size=9, show=True):
    import analyse
    import analyse.plotter
    import pylab
    import matplotlib
    import numpy

    def_font_size = matplotlib.rcParams['font.size']
    def_font_family = matplotlib.rcParams['font.family']
    matplotlib.rcParams['font.size'] = tick_labels_size
    if font_family is not None: matplotlib.rcParams['font.family'] = font_family

    if figure is None:
        figure = pylab.figure()
    else:
        show = False

    if type(subplot) is int:
        subplot = analyse.plotter.get_subplot(subplot)
    elif type(subplot) in [list, tuple]:
        if len(subplot) == 3:
            subplot = analyse.plotter.get_subplot(subplot)

    crval = raw.data['crval1'][0]
    crpix = raw.data['crpix1'][0]
    cdelt = raw.data['cdelt1'][0]
    num = len(raw.data['data'][0])
    v = crval+(numpy.arange(num)-crpix+1)*cdelt

    ax = figure.add_axes(subplot)
    #[ax.plot(v, s) for s in spectra]
    [ax.plot(v, s) for s in raw.data['DATA']]
    ax.grid(grid)
    ax.set_title(title)

    if show:
        pylab.show()

    matplotlib.rcParams['font.size'] = def_font_size
    matplotlib.rcParams['font.family'] = def_font_family
    return figure

def draw_ps_spectrum(cube, figure=None, subplot=111, title='', grid=True,
                      font_family=None, tick_labels_size=9, show=True, *args, **kwargs):
    import analyse
    import analyse.plotter
    import pylab
    import matplotlib

    def_font_size = matplotlib.rcParams['font.size']
    def_font_family = matplotlib.rcParams['font.family']
    matplotlib.rcParams['font.size'] = tick_labels_size
    if font_family is not None: matplotlib.rcParams['font.family'] = font_family

    if figure is None:
        figure = pylab.figure()
    else:
        show = False

    if type(subplot) is int:
        subplot = analyse.plotter.get_subplot(subplot)
    elif type(subplot) in [list, tuple]:
        if len(subplot) == 3:
            subplot = analyse.plotter.get_subplot(subplot)

    nz, ny, nx = cube.data.shape
    spectra = cube.data.T.reshape(nx*ny, nz)
    v = analyse.generate_axis(cube, axis=3) / 1000.

    ax = figure.add_axes(subplot)
    [ax.plot(v, s) for s in spectra]
    ax.grid(grid)
    ax.set_title(title)

    if show:
        pylab.show()

    matplotlib.rcParams['font.size'] = def_font_size
    matplotlib.rcParams['font.family'] = def_font_family
    return figure

def draw_ss_spectrum(cube, flag, figure=None, subplot=111, title='', grid=True,
                      font_family=None, tick_labels_size=9, show=True, *args, **kwargs):
    import analyse
    import analyse.plotter
    import pylab
    import numpy
    import matplotlib

    tpeak, rms = analyse_ss(cube, flag)

    def_font_size = matplotlib.rcParams['font.size']
    def_font_family = matplotlib.rcParams['font.family']
    matplotlib.rcParams['font.size'] = tick_labels_size
    if font_family is not None: matplotlib.rcParams['font.family'] = font_family

    if figure is None:
        figure = pylab.figure()
    else:
        show = False

    if type(subplot) is int:
        subplot = analyse.plotter.get_subplot(subplot)
    elif type(subplot) in [list, tuple]:
        if len(subplot) == 3:
            subplot = analyse.plotter.get_subplot(subplot)

    nz, ny, nx = cube.data.shape
    spectra = cube.data.T.reshape(nx*ny, nz)
    v = analyse.generate_axis(cube, axis=3) / 1000.

    ax = figure.add_axes(subplot)
    [ax.plot(v, s) for s in spectra]
    ax.grid(grid)
    ax.set_title(title)
    """
    if fits_params['MOLECULE']=='12CO': tmb = float(sslist['tmb12'])
    elif fits_params['MOLECULE']=='13CO': tmb = float(sslist['tmb13'])
    elif fits_params['MOLECULE']=='C18O': tmb = float(sslist['tmb18'])
    else: pass
    """
    # 判断の参考に参考にTmbが○％のところに線を引く
    """
    xs = float(sslist['v_center']) - float(sslist['v_width'])/2.
    xe = float(sslist['v_center']) + float(sslist['v_width'])/2.
    """
    ###### atodekesu #####
    xs = -41.       #
    xe = 55.        #
    tmb = 40.       #
    ###### atodekesu #####
    percents = numpy.arange(50, 101, 10)
    percent_lines = tmb * percents / 100.
    [pylab.axhline(y=_y, color='r', linestyle='--', linewidth=2, alpha=0.5) for _y in percent_lines]
    [ax.text(xs, _y, str(_p)+'%') for _y, _p in zip(percent_lines, percents)]
    ax.set_ylim(-2, tmb*1.2)
    ax.set_ylabel('Ta* [K]')
    ax.set_xlabel('km/s')
    ax.set_xlim(xs, xe)
    #ax.set_title('%s %s' % (fits_params['OBJECT'], fits_params['DATE-OBS']))
    ax.annotate('Tpeak=%.2f [K]\nTmb=%.2f [K]\neff.=%.2f\nTrms=%.2f [K]'%(tpeak, tmb, tpeak/tmb, rms), xy=(.65,.85), xycoords='axes fraction')

    """
    if str(fits_params['BACKEND'])[-1]==1: direction='H'
    elif str(fits_params['BACKEND'])[-1]==2: direction='V'
    else: pass
    """

    if show:
        pylab.show()

    matplotlib.rcParams['font.size'] = def_font_size
    matplotlib.rcParams['font.family'] = def_font_family
    return figure

def analyse_ss(cube, flag):
    import numpy

    g = lambda x, mu, sig: 1/numpy.sqrt(2*numpy.pi*sig**2) * numpy.exp(-(x-mu)**2/(2*sig**2.))
    x = numpy.arange(len(cube.data))
    mu = len(cube.data)/2
    sig = 20/(2*numpy.sqrt(2*numpy.log(2)))
    g0 = g(x, mu, sig)  #コンボリューション用のガウシアンを生成
    g0 /= numpy.sum(g0)  #規格化

    nz, ny, nx = cube.data.shape
    spectra = cube.data.copy().T.reshape(nx*ny, nz)
    spectrum = spectra[0]
    spectrum[numpy.isnan(spectrum)] = 0
    spectrum[numpy.isinf(spectrum)] = 0
    sd = numpy.convolve(spectrum, v=g0, mode='same')  #コンボリューション
    d_sd = [sd[i+1]-sd[i] for i,d_ in list(enumerate(sd))[:-1]]  #微分
    dd_sd = [d_sd[i+1]-d_sd[i] for i,d_ in list(enumerate(d_sd))[:-1]]  #微分
    # 観測された標準天体のピーク温度のチャンネルを算出(2回微分を行ったときの極小値)
    ddmin_ind = numpy.where(dd_sd==numpy.min(dd_sd))[0][0] - 1
    # 観測された標準天体のピーク温度を算出(上で求めたチャンネルの温度と、前後2点の値との平均をとっている)
    tpeak = numpy.average(spectrum[ddmin_ind-2:ddmin_ind+3])

    # --- rmsを求める ---
    flags = flag.data.copy().T.reshape(nx*ny, nz)
    spectra[numpy.where(flags!=0)] = 0
    spectra[numpy.isnan(spectra)] = 0
    spectra[numpy.isinf(spectra)] = 0
    rms = [numpy.sqrt(numpy.sum(spectra[i,:]**2)/float(nz)) for i in range(len(spectra))][0]

    print('tpeak: %.3f rms: %.3f'%(tpeak,rms))

    return tpeak, rms

def draw_ps(raw=None, cube=None, suptitle='', show=True,
             *args, **kwargs):
    import pylab

    fig = pylab.figure(figsize=(10,8))
    fig.suptitle(suptitle)

    #plot = custom_draw_map(figure=fig, show=False)
    draw_ps_raw(raw, figure=fig, subplot=231, title='raw', *args, **kwargs)
    draw_ps_spectrum(cube, figure=fig, subplot=234, title='spectra', *args, **kwargs)

    if show:
        pylab.show()

    return fig


import os
from PySide6.QtWidgets import QDialog, QFileSystemModel, QAbstractItemView, QComboBox
from PySide6.QtCore import Qt, Slot, QDir, QStandardPaths, QModelIndex
from PySide6.QtGui import QKeySequence, QValidator, QAction
from spinetoolbox.helpers import ProjectDirectoryIconProvider
from spinetoolbox.widgets.notification import Notification
from spinetoolbox.widgets.custom_menus import OpenProjectDialogComboBoxContextMenu

class OpenProjectDialog(QDialog):
    """A dialog that lets user select a project to open either by choosing
    an old .proj file or by choosing a project directory."""

    def __init__(self, toolbox):
        """

        Args:
            toolbox (ToolboxUI): QMainWindow instance
        """
        from ..ui import open_project_dialog  # pylint: disable=import-outside-toplevel

        super().__init__(parent=toolbox, f=Qt.Dialog)  # Setting the parent inherits the stylesheet
        self._toolbox = toolbox
        # Set up the user interface from Designer file
        self.ui = open_project_dialog.Ui_Dialog()
        self.ui.setupUi(self)
        self.combobox_context_menu = None
        # Ensure this dialog is garbage-collected when closed
        self.setAttribute(Qt.WA_DeleteOnClose)
        self.setWindowFlag(Qt.WindowType.WindowContextHelpButtonHint, False)
        # QActions for keyboard shortcuts
        self.go_root_action = QAction(self)
        self.go_home_action = QAction(self)
        self.go_documents_action = QAction(self)
        self.go_desktop_action = QAction(self)
        self.set_keyboard_shortcuts()
        self.selected_path = ""
        self.cb_ss = self.ui.comboBox_current_path.styleSheet()
        self.file_model = CustomQFileSystemModel()
        self.file_model.setFilter(QDir.AllDirs | QDir.NoDotAndDotDot)
        self.icon_provider = ProjectDirectoryIconProvider()
        self.file_model.setIconProvider(self.icon_provider)
        self.file_model.setRootPath(QDir.rootPath())
        self.ui.treeView_file_system.setModel(self.file_model)
        self.file_model.sort(0, Qt.AscendingOrder)
        # Enable validator (experimental, not very useful here)
        # Validator prevents typing Invalid strings to combobox. (not in use)
        # When text in combobox is Intermediate, the validator prevents emitting
        # currentIndexChanged signal when enter is pressed.
        # Pressing enter still triggers the done() slot of the QDialog.
        self.validator = DirValidator()
        self.ui.comboBox_current_path.setValidator(self.validator)
        self.ui.comboBox_current_path.setInsertPolicy(QComboBox.InsertPolicy.NoInsert)
        # Read recent project directories and populate combobox
        recents = self._toolbox.qsettings().value("appSettings/recentProjectStorages", defaultValue=None)
        if recents:
            recents_lst = str(recents).split("\n")
            self.ui.comboBox_current_path.insertItems(0, recents_lst)
            # Set start index to most recent project storage or to root if it does not exist
            p = self.ui.comboBox_current_path.itemText(0)
            if os.path.isdir(p):
                start_index = self.file_model.index(p)
            else:
                start_index = self.file_model.index(QDir.homePath())
        else:
            start_index = self.file_model.index(QDir.homePath())
            self.ui.comboBox_current_path.setCurrentIndex(-1)
        self.file_model.directoryLoaded.connect(self.expand_and_resize)
        # Start browsing to start index immediately when dialog is shown
        self.start_path = self.file_model.filePath(start_index)
        self.starting_up = True
        self.ui.treeView_file_system.setCurrentIndex(start_index)
        self.connect_signals()

    def set_keyboard_shortcuts(self):
        """Creates keyboard shortcuts for the 'Root', 'Home', etc. buttons."""
        self.go_root_action.setShortcut(QKeySequence(Qt.Key_F1))
        self.addAction(self.go_root_action)
        self.go_home_action.setShortcut(QKeySequence(Qt.Key_F2))
        self.addAction(self.go_home_action)
        self.go_documents_action.setShortcut(QKeySequence(Qt.Key_F3))
        self.addAction(self.go_documents_action)
        self.go_desktop_action.setShortcut(QKeySequence(Qt.Key_F4))
        self.addAction(self.go_desktop_action)

    def connect_signals(self):
        """Connects signals to slots."""
        self.ui.toolButton_root.clicked.connect(self.go_root)
        self.ui.toolButton_home.clicked.connect(self.go_home)
        self.ui.toolButton_documents.clicked.connect(self.go_documents)
        self.ui.toolButton_desktop.clicked.connect(self.go_desktop)
        self.ui.comboBox_current_path.editTextChanged.connect(self.combobox_text_edited)
        self.ui.comboBox_current_path.currentIndexChanged.connect(self.current_index_changed)
        self.ui.comboBox_current_path.customContextMenuRequested.connect(self.show_context_menu)
        self.validator.changed.connect(self.validator_state_changed)
        self.ui.treeView_file_system.clicked.connect(self.set_selected_path)
        self.ui.treeView_file_system.doubleClicked.connect(self.open_project)
        self.ui.treeView_file_system.selectionModel().currentChanged.connect(self.current_changed)
        self.go_root_action.triggered.connect(self.go_root)
        self.go_home_action.triggered.connect(self.go_home)
        self.go_documents_action.triggered.connect(self.go_documents)
        self.go_desktop_action.triggered.connect(self.go_desktop)

    @Slot(str)
    def expand_and_resize(self, p):
        """Expands, resizes, and scrolls the tree view to the current directory
        when the file model has finished loading the path. Slot for the file
        model's directoryLoaded signal. The directoryLoaded signal is emitted only
        if the directory has not been cached already. Note, that this is
        only used when the open project dialog is opened

        Args:
             p (str): Directory that has been loaded
        """
        if self.starting_up:
            current_index = self.ui.treeView_file_system.currentIndex()
            self.ui.treeView_file_system.scrollTo(current_index, hint=QAbstractItemView.PositionAtTop)
            self.ui.treeView_file_system.expand(current_index)
            if p == self.start_path:
                self.ui.treeView_file_system.resizeColumnToContents(0)
                self.set_selected_path(current_index)
                self.starting_up = False

    @Slot()
    def validator_state_changed(self):
        """Changes the combobox border color according to the current state of the validator."""
        state = self.ui.comboBox_current_path.validator().state
        if state == QValidator.State.Acceptable:
            self.ui.comboBox_current_path.setStyleSheet(self.cb_ss)
        elif state == QValidator.State.Intermediate:
            ss = "QComboBox {border: 1px solid #ff704d}"
            self.ui.comboBox_current_path.setStyleSheet(ss)
        else:  # Invalid. This is never returned (on purpose).
            ss = "QComboBox {border: 1px solid #ff3300}"
            self.ui.comboBox_current_path.setStyleSheet(ss)

    @Slot(int)
    def current_index_changed(self, i):
        """Combobox selection changed. This slot is processed when a new item
        is selected from the drop-down list. This is not processed when new
        item txt is QValidotor.Intermediate.

        Args:
            i (int): Selected row in combobox
        """
        p = self.ui.comboBox_current_path.itemText(i)
        if not os.path.isdir(p):
            self.remove_directory_from_recents(p, self._toolbox.qsettings())
            return
        fm_index = self.file_model.index(p)
        self.ui.treeView_file_system.collapseAll()
        self.ui.treeView_file_system.setCurrentIndex(fm_index)
        self.ui.treeView_file_system.expand(fm_index)
        self.ui.treeView_file_system.scrollTo(fm_index, hint=QAbstractItemView.PositionAtTop)

    @Slot("QModelIndex", "QModelIndex", name="current_changed")
    def current_changed(self, current, previous):
        """Processed when the current item in file system tree view has been
        changed with keyboard or mouse. Updates the text in combobox.

        Args:
            current (QModelIndex): Currently selected index
            previous (QModelIndex): Previously selected index
        """
        self.set_selected_path(current)

    @Slot("QModelIndex", name="set_selected_path")
    def set_selected_path(self, index):
        """Sets the text in the combobox as the selected path in the file system tree view.

        Args:
            index (QModelIndex): The index which was mouse clicked.
        """
        if not index.isValid():
            return
        selected_path = os.path.abspath(self.file_model.filePath(index))
        self.ui.comboBox_current_path.setCurrentText(selected_path)  # Emits editTextChanged signal
        self.selected_path = selected_path

    @Slot(str)
    def combobox_text_edited(self, text):
        """Updates selected path when combobox text is edited.
        Note: pressing enter in combobox does not trigger this.
        """
        self.selected_path = text

    def selection(self):
        """Returns the selected path from dialog."""
        return os.path.abspath(self.selected_path)

    @Slot(bool, name="go_root")
    def go_root(self, checked=False):
        """Slot for the 'Root' button. Scrolls the treeview to show and select the user's root directory.

        Note: We need to expand and scroll the tree view here after setCurrentIndex
        just in case the directory has been loaded already.
        """
        self.ui.comboBox_current_path.setCurrentIndex(-1)
        root_index = self.file_model.index(QDir.rootPath())
        self.ui.treeView_file_system.collapseAll()
        self.ui.treeView_file_system.setCurrentIndex(root_index)
        self.ui.treeView_file_system.expand(root_index)
        self.ui.treeView_file_system.scrollTo(root_index, hint=QAbstractItemView.PositionAtTop)

    @Slot(bool, name="go_home")
    def go_home(self, checked=False):
        """Slot for the 'Home' button. Scrolls the treeview to show and select the user's home directory."""
        self.ui.comboBox_current_path.setCurrentIndex(-1)
        home_index = self.file_model.index(QDir.homePath())
        self.ui.treeView_file_system.collapseAll()
        self.ui.treeView_file_system.setCurrentIndex(home_index)
        self.ui.treeView_file_system.expand(home_index)
        self.ui.treeView_file_system.scrollTo(home_index, hint=QAbstractItemView.PositionAtTop)

    @Slot(bool, name="go_documents")
    def go_documents(self, checked=False):
        """Slot for the 'Documents' button. Scrolls the treeview to show and select the user's documents directory."""
        docs = QStandardPaths.writableLocation(QStandardPaths.DocumentsLocation)
        if not docs:
            return
        self.ui.comboBox_current_path.setCurrentIndex(-1)
        docs_index = self.file_model.index(docs)
        self.ui.treeView_file_system.collapseAll()
        self.ui.treeView_file_system.setCurrentIndex(docs_index)
        self.ui.treeView_file_system.expand(docs_index)
        self.ui.treeView_file_system.scrollTo(docs_index, hint=QAbstractItemView.PositionAtTop)

    @Slot(bool, name="go_desktop")
    def go_desktop(self, checked=False):
        """Slot for the 'Desktop' button. Scrolls the treeview to show and select the user's desktop directory."""
        desktop = QStandardPaths.writableLocation(QStandardPaths.DesktopLocation)  # Return a list
        if not desktop:
            return
        self.ui.comboBox_current_path.setCurrentIndex(-1)
        desktop_index = self.file_model.index(desktop)
        self.ui.treeView_file_system.collapseAll()
        self.ui.treeView_file_system.setCurrentIndex(desktop_index)
        self.ui.treeView_file_system.expand(desktop_index)
        self.ui.treeView_file_system.scrollTo(desktop_index, hint=QAbstractItemView.PositionAtTop)

    @Slot("QModelIndex")
    def open_project(self, index):
        """Opens project if index contains a valid Spine Toolbox project.
        Slot for the mouse doubleClicked signal. Prevents showing the
        'Not a valid spine toolbox project' notification if user just wants
        to collapse a directory.

        Args:
            index (QModelIndex): File model index which was double clicked
        """
        if not index.isValid():
            return
        possible_project_json_file = os.path.join(self.selection(), ".spinetoolbox", "project.json")
        if not os.path.isfile(possible_project_json_file):
            return
        self.done(QDialog.DialogCode.Accepted)

    def done(self, r):
        """Checks that selected path exists and is a valid
        Spine Toolbox directory when ok button is clicked or
        when enter is pressed without the combobox being in focus.

        Args:
            r (int) Return code
        """
        if r == QDialog.DialogCode.Accepted:
            if not os.path.isdir(self.selection()):
                notification = Notification(self, "Path does not exist")
                notification.show()
                return
            project_json_fp = os.path.abspath(os.path.join(self.selection(), ".spinetoolbox", "project.json"))
            if not os.path.isfile(project_json_fp):
                notification = Notification(self, "Not a valid Spine Toolbox project")
                notification.show()
                return
            # self.selection() now contains a valid Spine Toolbox project directory.
            # Add the parent directory of selected directory to qsettings
            self.update_recents(
                os.path.abspath(os.path.join(self.selection(), os.path.pardir)), self._toolbox.qsettings()
            )
        super().done(r)

    @staticmethod
    def update_recents(entry, qsettings):
        """Adds a new entry to QSettings variable that remembers the five most recent project storages.

        Args:
            entry (str): Abs. path to a directory that most likely contains other Spine Toolbox Projects as well.
                First entry is also used as the initial path for File->New Project dialog.
            qsettings (QSettings): Toolbox qsettings object
        """
        recents = qsettings.value("appSettings/recentProjectStorages", defaultValue=None)
        if not recents:
            updated_recents = entry
        else:
            recents = str(recents)
            recents_list = recents.split("\n")
            # Add path only if it's not in the list already
            if entry not in recents_list:
                recents_list.insert(0, entry)
                if len(recents_list) > 5:
                    recents_list.pop()
            else:
                # If entry was on the list, move it as the first item
                recents_list.insert(0, recents_list.pop(recents_list.index(entry)))
            updated_recents = "\n".join(recents_list)
        # Save updated recent paths
        qsettings.setValue("appSettings/recentProjectStorages", updated_recents)
        qsettings.sync()  # Commit change immediately

    @staticmethod
    def remove_directory_from_recents(p, qsettings):
        """Removes directory from the recent project storages.

        Args:
            p (str): Full path to a project directory
            qsettings (QSettings): Toolbox qsettings object
        """
        recents = qsettings.value("appSettings/recentProjectStorages", defaultValue=None)
        if not recents:
            return
        recents = str(recents)
        recents_list = recents.split("\n")
        if p in recents_list:
            recents_list.pop(recents_list.index(p))
        updated_recents = "\n".join(recents_list)
        # Save updated recent paths
        qsettings.setValue("appSettings/recentProjectStorages", updated_recents)
        qsettings.sync()  # Commit change immediately

    @Slot("QPoint")
    def show_context_menu(self, pos):
        """Shows the context menu for the QCombobox with a 'Clear history' entry.

        Args:
            pos (QPoint): Mouse position
        """
        # ind = self.ui.comboBox_current_path.indexAt(pos)
        global_pos = self.ui.comboBox_current_path.mapToGlobal(pos)
        # if not ind.isValid():
        self.combobox_context_menu = OpenProjectDialogComboBoxContextMenu(self, global_pos)
        action = self.combobox_context_menu.get_action()
        if action == "Clear history":
            self.ui.comboBox_current_path.clear()
            self._toolbox.qsettings().setValue("appSettings/recentProjectStorages", "")
            self.go_root()
        else:  # No option selected
            pass
        self.combobox_context_menu.deleteLater()
        self.combobox_context_menu = None

    def closeEvent(self, event=None):
        """Handles dialog closing.

        Args:
            event (QCloseEvent): Close event
        """
        if event:
            event.accept()

class CustomQFileSystemModel(QFileSystemModel):
    """Custom file system model."""

    def columnCount(self, parent=QModelIndex()):
        """Returns one."""
        return 1

class DirValidator(QValidator):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.state = None

    def validate(self, txt, pos):
        """Returns Invalid if input is invalid according to this
        validator's rules, Intermediate if it is likely that a
        little more editing will make the input acceptable and
        Acceptable if the input is valid.

        Args:
            txt (str): Text to validate
            pos (int): Cursor position

        Returns:
            QValidator.State: Invalid, Intermediate, or Acceptable
        """
        previous_state = self.state
        if not txt:
            self.state = QValidator.State.Intermediate
            if not previous_state == self.state:
                self.changed.emit()
            return self.state
        if os.path.isdir(txt):
            self.state = QValidator.State.Acceptable
            if not previous_state == self.state:
                self.changed.emit()
            return self.state
        self.state = QValidator.State.Intermediate
        if not previous_state == self.state:
            self.changed.emit()
        return self.state


import time
import math
from copy import copy, deepcopy
import os
import os.path as osp
from collections import defaultdict
import signal

import wandb
import numpy as np

import logger
# from helpers.distributed_util import sync_check
from helpers.console_util import timed_cm_wrapper, log_iter_info
from helpers.opencv_util import record_video

BRAC_BEHAV_STEPS = 25_000  # note: this has nothing to do with 'warm-starting with behavioral cloning'

debug_lvl = os.environ.get('DEBUG_LVL', 0)
try:
    debug_lvl = np.clip(int(debug_lvl), a_min=0, a_max=3)
except ValueError:
    debug_lvl = 0
DEBUG = bool(debug_lvl >= 1)

def rollout_generator(env, agent, rollout_len, use_noise_process):

    t = 0
    if use_noise_process:
        # Reset agent's noise process
        agent.reset_noise()
    # Reset agent's env
    ob = np.array(env.reset())

    while True:

        # Predict action
        ac = agent.predict(ob, apply_noise=True, which='main')
        # NaN-proof and clip
        ac = np.nan_to_num(ac)
        ac = np.clip(ac, env.action_space.low, env.action_space.high)

        if t > 0 and t % rollout_len == 0:
            yield

        # Interact with env(s)
        new_ob, rew, done, _ = env.step(ac)

        # Assemble and store transition in memory
        transition = {
            "obs0": ob,
            "acs": ac,
            "obs1": new_ob,
            "rews": rew,
            "dones1": done,
        }
        agent.store_transition(transition)

        # Set current state with the next
        ob = np.array(deepcopy(new_ob))

        if done:
            if use_noise_process:
                # Reset agent's noise process
                agent.reset_noise()
            # Reset agent's env
            ob = np.array(env.reset())

        t += 1

def ep_generator(env, agent, render, record, which):
    """Generator that spits out a trajectory collected during a single episode
    `append` operation is also significantly faster on lists than numpy arrays,
    they will be converted to numpy arrays once complete and ready to be yielded.
    """

    kwargs = {'mode': 'rgb_array'}

    def _render():
        return env.render(**kwargs)

    ob = np.array(env.reset())

    if record:
        ob_orig = _render()

    cur_ep_len = 0
    cur_ep_env_ret = 0
    obs = []
    if record:
        obs_render = []
    acs = []
    env_rews = []

    # Not necessarily used
    qs = []
    mcqs = []

    while True:

        # Predict action
        _ac = agent.predict(ob, apply_noise=False, which=which)

        if isinstance(_ac, dict):
            ac = _ac['ac']
            q = _ac['q_mean']
            mcq = _ac['mc_q_mean']
            qs.append(q)
            mcqs.append(mcq)
        else:
            ac = _ac

        # NaN-proof and clip
        ac = np.nan_to_num(ac)
        ac = np.clip(ac, env.action_space.low, env.action_space.high)

        obs.append(ob)
        if record:
            obs_render.append(ob_orig)
        acs.append(ac)
        new_ob, env_rew, done, _ = env.step(ac)

        if render:
            env.render()

        if record:
            ob_orig = _render()

        env_rews.append(env_rew)
        cur_ep_len += 1
        cur_ep_env_ret += env_rew
        ob = np.array(deepcopy(new_ob))
        if done:
            obs = np.array(obs)
            if record:
                obs_render = np.array(obs_render)
            acs = np.array(acs)
            env_rews = np.array(env_rews)
            out = {"obs": obs,
                   "acs": acs,
                   "env_rews": env_rews,
                   "ep_len": cur_ep_len,
                   "ep_env_ret": cur_ep_env_ret}
            if isinstance(_ac, dict):
                out.update({'q': np.array(qs).mean(),
                            'mcq': np.array(mcqs).mean()})
            if record:
                out.update({"obs_render": obs_render})

            yield out

            cur_ep_len = 0
            cur_ep_env_ret = 0
            obs = []
            if record:
                obs_render = []
            acs = []
            env_rews = []

            # Not necessarily used
            qs = []
            mcqs = []

            ob = np.array(env.reset())

            if record:
                ob_orig = _render()

def evaluate(args,
             env,
             agent_wrapper,
             experiment_name):

    # Create an agent
    agent = agent_wrapper()

    # Create episode generator
    ep_gen = ep_generator(env, agent, args.render)

    if args.record:
        vid_dir = osp.join(args.video_dir, experiment_name)
        os.makedirs(vid_dir, exist_ok=True)

    # Load the model
    agent.load(args.model_path, args.iter_num)
    logger.info("model loaded from path:\n  {}".format(args.model_path))

    # Initialize the history data structures
    ep_lens = []
    ep_env_rets = []
    # Collect trajectories
    for i in range(args.num_trajs):
        logger.info("evaluating [{}/{}]".format(i + 1, args.num_trajs))
        traj = ep_gen.__next__()
        ep_len, ep_env_ret = traj['ep_len'], traj['ep_env_ret']
        # Aggregate to the history data structures
        ep_lens.append(ep_len)
        ep_env_rets.append(ep_env_ret)
        if args.record:
            # Record a video of the episode
            record_video(vid_dir, i, traj['obs_render'])

    # Log some statistics of the collected trajectories
    ep_len_mean = np.mean(ep_lens)
    ep_env_ret_mean = np.mean(ep_env_rets)
    logger.record_tabular("ep_len_mean", ep_len_mean)
    logger.record_tabular("ep_env_ret_mean", ep_env_ret_mean)
    logger.dump_tabular()

def generate_buffer(args,
                    env,
                    agent_wrapper,
                    experiment_name):

    # Set up where to save the exploration data
    explo_dir = osp.join(args.exploration_dir, experiment_name)
    os.makedirs(explo_dir, exist_ok=True)

    # Create an agent
    agent = agent_wrapper()

    # Log the type of exploration used by the agent
    logger.info("parameter noise used: {}".format(agent.param_noise))
    logger.info("action noise used: {}".format(agent.ac_noise))

    # Create episode generator
    roll_gen = rollout_generator(env, agent, rollout_len=args.gen_buffer_size)
    # Note, we only need to generate once from the generator (no training).

    # Load the model
    agent.load(args.model_path, args.iter_num)
    logger.info("model loaded from path:\n  {}".format(args.model_path))

    # Generate rollout to populate the replay buffer
    roll_gen.__next__()  # no need to get the returned rollout, stored in buffer

    print(agent.replay_buffer.num_entries)

    # Save the data on disk
    agent.save_full_history(explo_dir, "0-generated")  # technically no training is done here
    # Note, we use the full history method to be sure nothing is flushed

def learn(args,
          rank,
          env,
          main_eval_env,
          maxq_eval_env,
          cwpq_eval_env,
          agent_wrapper,
          experiment_name,
          use_noise_process=None):

    # Create an agent
    agent = agent_wrapper()

    # Create context manager that records the time taken by encapsulated ops
    timed = timed_cm_wrapper(logger, use=DEBUG)

    # Start clocks
    num_iters = int(args.num_steps) if args.offline else int(args.num_steps) // args.rollout_len
    iters_so_far = 0
    timesteps_so_far = 0
    tstart = time.time()

    if rank == 0:
        # Create collections
        d = defaultdict(list)

        # Set up model save directory
        ckpt_dir = osp.join(args.checkpoint_dir, experiment_name)
        os.makedirs(ckpt_dir, exist_ok=True)
        if args.record:
            vid_dir = osp.join(args.video_dir, experiment_name)
            os.makedirs(vid_dir, exist_ok=True)
        if not agent.hps.offline:
            # Set up replay save directory
            replay_dir = osp.join(args.replay_dir, experiment_name)
            os.makedirs(replay_dir, exist_ok=True)

        # Handle timeout signal gracefully
        def timeout(signum, frame):
            # Save the model
            agent.save(ckpt_dir, "{}_timeout".format(iters_so_far))
            # No need to log a message, orterun stopped the trace already
            # No need to end the run by hand, SIGKILL is sent by orterun fast enough after SIGTERM

        # Tie the timeout handler with the termination signal
        # Note, orterun relays SIGTERM and SIGINT to the workers as SIGTERM signals,
        # quickly followed by a SIGKILL signal (Open-MPI impl)
        signal.signal(signal.SIGTERM, timeout)

        # Group by everything except the seed, which is last, hence index -1
        group = '.'.join(experiment_name.split('.')[:-1])

        # Set up wandb
        while True:
            try:
                wandb.init(
                    project=args.wandb_project,
                    name=experiment_name,
                    id=experiment_name,
                    group=group,
                    config=args.__dict__,
                    dir=args.root,
                )
                break
            except Exception:
                pause = 10
                logger.info("wandb co error. Retrying in {} secs.".format(pause))
                time.sleep(pause)
        logger.info("wandb co established!")

    # Create rollout generator for training the agent
    if args.offline:
        pass
    else:
        roll_gen = rollout_generator(env, agent, args.rollout_len, use_noise_process)
    if rank == 0:
        # Create episode generator for evaluating the agent
        main_eval_ep_gen = ep_generator(main_eval_env, agent, args.render, args.record, which='main')
        maxq_eval_ep_gen = ep_generator(maxq_eval_env, agent, args.render, args.record, which='maxq')
        cwpq_eval_ep_gen = ep_generator(cwpq_eval_env, agent, args.render, args.record, which='cwpq')

    while iters_so_far <= num_iters:

        if iters_so_far % 100 == 0 or DEBUG:
            log_iter_info(logger, iters_so_far, num_iters, tstart)

        # if iters_so_far % 20 == 0:
        #     # Check if the mpi workers are still synced
        #     sync_check(agent.actr)

        if rank == 0 and iters_so_far % args.save_frequency == 0:
            # Save the model
            agent.save(ckpt_dir, iters_so_far)
            logger.info("time to checkpoint. Saving model @: {}".format(ckpt_dir))
            if not agent.hps.offline:
                # Save the replay buffer and the full history
                agent.save_memory(replay_dir, iters_so_far)
                agent.save_full_history(replay_dir, iters_so_far)

        if args.offline:
            pass
        else:
            # Sample mini-batch in env with perturbed actor and store transitions
            with timed("interacting"):
                roll_gen.__next__()  # no need to get the returned rollout, stored in buffer

        if args.algo.split('_')[0] == 'brac' and iters_so_far == 0:
            with timed('behavior policy pre-training'):
                # Train the behavior policy estimate in a pre-training phase, as in the BRAC codebase
                for i in range(BRAC_BEHAV_STEPS):
                    # Sample a batch of transitions from the replay buffer
                    batch = agent.sample_batch()
                    loss = agent.update_behavior_policy_clone(batch=batch)
                    if i % 10 == 0:
                        i_fmt = str(i).zfill(math.floor(math.log10(BRAC_BEHAV_STEPS)) + 1)
                        logger.info(f"bc iter: {i_fmt} | bc loss: {loss}")

        with timed('training'):
            for training_step in range(args.training_steps_per_iter):

                if agent.param_noise is not None:
                    if training_step % args.pn_adapt_frequency == 0:
                        # Adapt parameter noise
                        agent.adapt_param_noise()
                    if rank == 0 and iters_so_far % args.eval_frequency == 0:
                        # Store the action-space dist between perturbed and non-perturbed
                        d['pn_dist'].append(agent.pn_dist)
                        # Store the new std resulting from the adaption
                        d['pn_cur_std'].append(agent.param_noise.cur_std)

                # Sample a batch of transitions from the replay buffer
                batch = agent.sample_batch()

                # Update the actor and critic
                metrics, lrnows = agent.update_actor_critic(
                    batch=batch,
                    update_actor=not bool(iters_so_far % args.actor_update_delay),  # from TD3
                    iters_so_far=iters_so_far,
                )
                if rank == 0 and iters_so_far % args.eval_frequency == 0:
                    # Log training stats
                    d['actr_losses'].append(metrics['actr_loss'])
                    d['crit_losses'].append(metrics['crit_loss'])
                    if iters_so_far >= agent.hps.warm_start:
                        if args.algo.split('_')[0] == 'bcp':
                            if agent.hps.base_pi_loss == 'crr_exp':
                                d['is_crr_adv_clipped_sum'].append(metrics['is_crr_adv_clipped_sum'])
                            d['alpha_pri'].append(metrics['alpha_pri'])
                            d['gap'].append(metrics['gap'])
                            if (agent.hps.use_rnd_monitoring or
                                    agent.hps.base_pe_loss in ['htg_1', 'htg_2'] or
                                    agent.hps.base_next_action == 'beta_theta_max_rnd'):
                                d['rnd_score'].append(metrics['rnd_score'])
                    if agent.hps.clipped_double:
                        d['twin_losses'].append(metrics['twin_loss'])
                    if agent.hps.prioritized_replay:
                        iws = metrics['iws']  # last one only

        if rank == 0 and iters_so_far % args.eval_frequency == 0:

            with timed("evaluating"):

                # Update the parameters of the eval actors with the ones of the trained one
                agent.update_eval_nets()

                for eval_step in range(args.eval_steps_per_iter):
                    # Sample an episode
                    main_eval_ep = main_eval_ep_gen.__next__()
                    maxq_eval_ep = maxq_eval_ep_gen.__next__()
                    cwpq_eval_ep = cwpq_eval_ep_gen.__next__()
                    # Aggregate data collected during the evaluation to the buffers
                    d['main_eval_len'].append(main_eval_ep['ep_len'])
                    d['maxq_eval_len'].append(maxq_eval_ep['ep_len'])
                    d['cwpq_eval_len'].append(cwpq_eval_ep['ep_len'])
                    d['main_eval_env_ret'].append(main_eval_ep['ep_env_ret'])
                    d['maxq_eval_env_ret'].append(maxq_eval_ep['ep_env_ret'])
                    d['cwpq_eval_env_ret'].append(cwpq_eval_ep['ep_env_ret'])
                    if args.algo.split('_')[0] == 'bcp':
                        d['eval_q'].append(maxq_eval_ep['q'])
                        d['eval_mcq'].append(maxq_eval_ep['mcq'])

        # Increment counters
        iters_so_far += 1
        if args.offline:
            _str = 'iters so far'
            step = copy(iters_so_far)
        else:
            timesteps_so_far += args.rollout_len
            _str = 'timesteps so far'
            step = copy(timesteps_so_far)

        if rank == 0 and ((iters_so_far - 1) % args.eval_frequency == 0):

            # Log stats in csv
            logger.record_tabular(_str, step)

            if iters_so_far - 1 >= agent.hps.warm_start:
                if args.algo.split('_')[0] == 'bcp':
                    if agent.hps.base_pi_loss == 'crr_exp':
                        logger.record_tabular('is_crr_adv_clipped_sum', np.mean(d['is_crr_adv_clipped_sum']))
                    logger.record_tabular('alpha_pri', np.mean(d['alpha_pri']))
                    logger.record_tabular('gap', np.mean(d['gap']))
                    if (agent.hps.use_rnd_monitoring or
                            agent.hps.base_pe_loss in ['htg_1', 'htg_2'] or
                            agent.hps.base_next_action == 'beta_theta_max_rnd'):
                        logger.record_tabular('rnd_score', np.mean(d['rnd_score']))
            logger.record_tabular('main_eval_len', np.mean(d['main_eval_len']))
            logger.record_tabular('maxq_eval_len', np.mean(d['maxq_eval_len']))
            logger.record_tabular('cwpq_eval_len', np.mean(d['cwpq_eval_len']))
            logger.record_tabular('main_eval_env_ret', np.mean(d['main_eval_env_ret']))
            logger.record_tabular('maxq_eval_env_ret', np.mean(d['maxq_eval_env_ret']))
            logger.record_tabular('cwpq_eval_env_ret', np.mean(d['cwpq_eval_env_ret']))
            if args.algo.split('_')[0] == 'bcp':
                logger.record_tabular('eval_q', np.mean(d['eval_q']))
                logger.record_tabular('eval_mcq', np.mean(d['eval_mcq']))
            logger.info("dumping stats in .csv file")
            logger.dump_tabular()

            if args.record:
                # Record the last episode in a video
                record_video(vid_dir, iters_so_far, maxq_eval_ep['obs_render'])  # picked only one, maxq

            # Log stats in dashboard
            if agent.hps.prioritized_replay:
                quantiles = [0.1, 0.25, 0.5, 0.75, 0.9]
                np.quantile(iws, quantiles)
                wandb.log({"q{}".format(q): np.quantile(iws, q)
                           for q in [0.1, 0.25, 0.5, 0.75, 0.9]},
                          step=step)
            if agent.param_noise is not None:
                wandb.log({'pn_dist': np.mean(d['pn_dist']),
                           'pn_cur_std': np.mean(d['pn_cur_std'])},
                          step=step)
            wandb.log({'actr_loss': np.mean(d['actr_losses']),
                       'actr_lrnow': np.array(lrnows['actr']),
                       'crit_loss': np.mean(d['crit_losses'])},
                      step=step)
            if agent.hps.clipped_double:
                wandb.log({'twin_loss': np.mean(d['twin_losses'])},
                          step=step)
            if iters_so_far - 1 >= agent.hps.warm_start:
                if args.algo.split('_')[0] == 'bcp':
                    if agent.hps.base_pi_loss == 'crr_exp':
                        wandb.log({'is_crr_adv_clipped_sum': np.mean(d['is_crr_adv_clipped_sum'])},
                                  step=step)
                    wandb.log({'alpha_pri': np.mean(d['alpha_pri']),
                               'gap': np.mean(d['gap'])},
                              step=step)
                    if (agent.hps.use_rnd_monitoring or
                            agent.hps.base_pe_loss in ['htg_1', 'htg_2'] or
                            agent.hps.base_next_action == 'beta_theta_max_rnd'):
                        wandb.log({'rnd_score': np.mean(d['rnd_score'])},
                                  step=step)
            wandb.log({'main_eval_len': np.mean(d['main_eval_len']),
                       'maxq_eval_len': np.mean(d['maxq_eval_len']),
                       'cwpq_eval_len': np.mean(d['cwpq_eval_len']),
                       'main_eval_env_ret': np.mean(d['main_eval_env_ret']),
                       'maxq_eval_env_ret': np.mean(d['maxq_eval_env_ret']),
                       'cwpq_eval_env_ret': np.mean(d['cwpq_eval_env_ret'])},
                      step=step)
            if args.algo.split('_')[0] == 'bcp':
                wandb.log({'eval_q': np.mean(d['eval_q']),
                           'eval_mcq': np.mean(d['eval_mcq'])},
                          step=step)

            # Clear the iteration's running stats
            d.clear()

    if not agent.hps.offline:
        # When running the algorithm online, save behavior policy model once we are done iterating
        if rank == 0:
            # Save the model
            agent.save(ckpt_dir, iters_so_far)
            logger.info("we're done. Saving behavior policy model @: {}".format(ckpt_dir))
            logger.info("bye.")
            # Save the replay buffer and the full history
            agent.save_memory(replay_dir, iters_so_far)
            agent.save_full_history(replay_dir, iters_so_far)

    if rank == 0:
        # Save once we are done iterating
        agent.save(ckpt_dir, iters_so_far)
        logger.info("we're done. Saving model @: {}".format(ckpt_dir))
        logger.info("bye.")


import textwrap

from PIL import Image
from PIL import ImageDraw
from PIL import ImageFont
from PIL import ImageFilter
from PIL import ImageChops

from .image import image
from .image import PImage
from .structure import push_style

from . import p5

__all__ = ['create_font', 'load_font', 'text', 'text_font',
    'text_align', 'text_leading', 'text_size', 'text_width',
    'text_ascent', 'text_descent'
    ]

_font_family = ImageFont.load_default()
_text_align_x = "LEFT"
_text_align_y = "TOP"
_text_leading = 0

def create_font(name, size=10):
    """Create the given font at the appropriate size.

    :param name: Filename of the font file (only pil, otf and ttf
        fonts are supported.)
    :type name: str

    :param size: Font size (only required when `name` refers to a
        truetype font; defaults to None)
    :type size: int | None

    """

    if name.endswith('ttf') or name.endswith('otf'):
        font = ImageFont.truetype(name, size)
    elif name.endswith('pil'):
        font = ImageFont.load(name)
    else:
        raise NotImplementedError("Font type not supported.")

    return font

def load_font(font_name):
    """Loads the given font into a font object

    """
    return create_font(font_name)

def text(*args, wrap_at=None):
    """Draw the given text on the screen and save the image.

    :param text_string: text to display
    :type text_string: str

    :param x: x-coordinate of text
    :type x: float

    :param y: y-coordinate of text
    :type y: float

    :param z: z-coordinate of text
    :type z: float

    :param position: position of the text on the screen
    :type position: tuple

    :param wrap_at: specifies the text wrapping column (defaults to
        None)
    :type wrap_at: int

    :returns: actual text that was drawn to the image (when wrapping
        is not set, this is just the unmodified text_string)
    :rtype: str

    """
    if len(args) == 2:
        text_string, position = args
    elif len(args) == 3 or len(args) == 4:
        text_string, position = args[0], args[1:]
    else:
        raise ValueError("Unexpected number of arguments passed to text()")

    if len(text_string) == 0:
        return

    global _font_family, _text_leading

    multiline = False
    if not (wrap_at is None):
        text_string = textwrap.fill(text_string, wrap_at)
        size = _font_family.getsize_multiline(text_string)
        multiline = True
    elif "\n" in text_string:
        multiline = True
        size = list(_font_family.getsize_multiline(text_string))
        size[1] += _text_leading*text_string.count("\n")
    else:
        size = _font_family.getsize(text_string)

    is_stroke_valid = False # True when stroke_weight != 0
    is_min_filter = False   # True when stroke_weight <0
    if p5.renderer.style.stroke_enabled:
        stroke_weight = p5.renderer.stroke_weight
        if stroke_weight < 0:
            stroke_weight = abs(stroke_weight)
            is_min_filter = True

        if stroke_weight > 0:
            if stroke_weight % 2 == 0:
                stroke_weight += 1
            is_stroke_valid = True

    if is_stroke_valid:
        new_size = list(map(lambda x:x+2*stroke_weight, size))
        is_stroke_valid = True
        text_xy = (stroke_weight, stroke_weight)
    else:
        new_size = size
        text_xy = (0,0)

    canvas = Image.new("RGBA", new_size, color=(0, 0, 0, 0))
    canvas_draw = ImageDraw.Draw(canvas)

    if multiline:
        canvas_draw.multiline_text(text_xy, text_string, font=_font_family, spacing=_text_leading)
    else:
        canvas_draw.text(text_xy, text_string, font=_font_family)

    text_image = PImage(*new_size)
    text_image._img = canvas

    if is_stroke_valid:
        if is_min_filter:
            canvas_dilate = canvas.filter(ImageFilter.MinFilter(stroke_weight))
        else:
            canvas_dilate = canvas.filter(ImageFilter.MaxFilter(stroke_weight))
        canvas_stroke = ImageChops.difference(canvas,canvas_dilate)
        text_stroke_image = PImage(*new_size)
        text_stroke_image._img = canvas_stroke

    width, height = new_size
    position = list(position)
    if _text_align_x == "LEFT":
        position[0] += 0
    elif _text_align_x == "RIGHT":
        position[0] -= width
    elif _text_align_x == "CENTER":
        position[0] -= width/2

    if _text_align_y == "TOP":
        position[1] += 0
    elif _text_align_y == "BOTTOM":
        position[1] -= height
    elif _text_align_y == "CENTER":
        position[1] -= height/2

    with push_style():
        if p5.renderer.style.fill_enabled:
            p5.renderer.style.tint_enabled = True
            p5.renderer.style.tint_color = p5.renderer.style.fill_color
            image(text_image, position)
        if p5.renderer.style.stroke_enabled and is_stroke_valid:
            p5.renderer.style.tint_enabled = True
            p5.renderer.style.tint_color = p5.renderer.style.stroke_color
            image(text_stroke_image, position)

    return text_string

def text_font(font, size=10):
    """Set current text font.

    :param font:
    :type font: PIL.ImageFont.ImageFont

    """
    global _font_family
    _font_family = font

def text_align(align_x, align_y=None):
    """Set the alignment of drawing text

    :param align_x: "RIGHT", "CENTER" or "LEFT".
    :type align_x: string

    :param align_y: "TOP", "CENTER" or "BOTTOM".
    :type align_y: string

    """

    global _text_align_x, _text_align_y
    _text_align_x = align_x

    if align_y:
        _text_align_y = align_y

def text_leading(leading):
    """Sets the spacing between lines of text in units of pixels

    :param leading: the size in pixels for spacing between lines
    :type align_x: int

    """

    global _text_leading
    _text_leading = leading

def text_size(size):
    """Sets the current font size

    :param leading: the size of the letters in units of pixels
    :type align_x: int

    """

    global _font_family

    # reload the font with new size
    if hasattr(_font_family, 'path'):
        if _font_family.path.endswith('ttf') or _font_family.path.endswith('otf'):
            _font_family = ImageFont.truetype(_font_family.path, size)
    else:
        raise ValueError("text_size is not supported for Bitmap Fonts")

def text_width(text):
    """Calculates and returns the width of any character or text string

    :param text_string: text
    :type text_string: str

    :returns: width of any character or text string
    :rtype: int

    """

    return _font_family.getsize(text)[0]

def text_ascent():
    """Returns ascent of the current font at its current size

    :returns: ascent of the current font at its current size
    :rtype: float

    """
    global _font_family
    ascent, descent = _font_family.getmetrics()
    return ascent

def text_descent():
    """Returns descent of the current font at its current size

    :returns:  descent of the current font at its current size
    :rtype: float

    """

    global _font_family
    ascent, descent = _font_family.getmetrics()
    return descent


# -*- coding: utf-8 -*-
"""
Created on Tue Sep  8 17:16:18 2020
@author: Cleiton Moya de Almeida

"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib
from sklearn.tree import DecisionTreeClassifier, export_graphviz
from sklearn.model_selection import  KFold, GridSearchCV
import graphviz

# Dataset reading
df = pd.read_table('banana.dat', skiprows=7, sep=',', names=['x1', 'x2', 'y'])
df = df.astype({'y': 'int32'})
X = df[['x1', 'x2']].to_numpy()
y = df.y.to_numpy()

param_grid = [dict(criterion = ['gini', 'entropy'],
                   min_impurity_decrease = [0, 0.01, 0.05, 0.1],
                   max_depth = [3, 4, 5],
                   max_leaf_nodes = [3, 5])]

# Cross-validation
K = 10 # k-fold parameter
cv = KFold(n_splits=K)
grid = GridSearchCV(DecisionTreeClassifier(), param_grid=param_grid, cv=cv, n_jobs=-1,verbose=10)
grid.fit(X,y)
best_std_score = grid.cv_results_['std_test_score'][grid.best_index_]

# Results:
print("Best parameters: %s \nAccuracy: %0.3f \u00B1 %0.3f"
      % (grid.best_params_, grid.best_score_, best_std_score))

# Model
clf = DecisionTreeClassifier(min_impurity_decrease=0.01,
                             max_depth = 3,
                             min_samples_leaf = 1,
                             max_leaf_nodes = 3
                             )
clf = clf.fit(X,y)

dot_data = export_graphviz(clf, out_file=None,
                           feature_names=['x1', 'x2'],
                           class_names=['-1', '+1'],
                           filled=True, rounded=True,
                           special_characters=True)

graph = graphviz.Source(dot_data)
graph.render("exp1-tree")

# Decision surface
plt.rcParams.update({'font.size': 8})
plt.figure(figsize=(4,3))
plot_step = 0.02
x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
xx, yy = np.meshgrid(np.arange(x_min, x_max, plot_step),
                     np.arange(y_min, y_max, plot_step))
plt.tight_layout(h_pad=0.5, w_pad=0.5, pad=2.5)

plot_colors = ['C0','C1']
cmap = matplotlib.colors.ListedColormap(plot_colors)
Z = clf.predict(np.c_[xx.ravel(), yy.ravel()])
Z = Z.reshape(xx.shape)
cs = plt.contourf(xx, yy, Z, cmap=cmap,alpha=0.5)

labels=['-1', '+1']

n_classes = 2
for i, color in zip([-1,1], plot_colors):
        idx = np.where(y == i)
        plt.scatter(X[idx, 0], X[idx, 1], c=color, label=labels[i],
                    cmap=plt.cm.RdYlBu, edgecolor='black', s=10,linewidths=0.2)

plt.xlim(-3, 3)
plt.ylim(-3, 3)
plt.xlabel('x1')
plt.ylabel('x2')
plt.tight_layout()


"""Contains application configurator.

Used to transfer parameters from user to application.
"""

import os
import re
import configparser

path = os.path.abspath(os.path.expanduser('~/.rapo/rapo.ini'))
encoding = 'utf-8'

class Configurator(dict):
    """Represents main configurator."""

    def __init__(self):
        super().__init__()
        if self.found:
            names = ['SCHEDULER', 'DATABASE', 'LOGGING', 'API']
            for name in names:
                self[name] = Configuration(name)
            self.load()

    def load(self):
        """Load configuration from file into memory."""
        parser = configparser.ConfigParser(allow_no_value=True)
        parser.read(path, encoding=encoding)
        for section in parser.sections():
            for option in parser.options(section):
                initial_value = parser[section][option]
                final_value = self.normalize(initial_value)
                self[section][option] = final_value
        return self

    def normalize(self, value):
        """Normalize given parameter value."""
        if value is None or value.upper() == 'NONE' or value.isspace():
            return None
        elif value.upper() == 'TRUE':
            return True
        elif value.upper() == 'FALSE':
            return False
        elif re.match(r'^[+-]?\d+$', value):
            return int(value)
        elif re.match(r'^[+-]?(\d*.\d+|\d+\.\d*)$', value):
            return float(value)
        return value

    def check(self, configuration_name):
        """Determine whether configuration presented or not."""
        if self.get(configuration_name) is not None:
            return True
        return False

    @property
    def found(self):
        """Determine whether configuration file found or not."""
        if os.path.exists(path):
            return True
        return False

class Configuration(dict):
    """Represents some configuration."""

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

    def __setitem__(self, key, value):
        """Set the parameter value."""
        super().__setitem__(key.lower(), value)

    def __getitem__(self, key):
        """Get the parameter value or raise an exception if not found."""
        return super().__getitem__(key.lower())

    def get(self, key, default=None):
        """Get the parameter value or return default if not found."""
        return super().get(key.lower(), default)

    def get_deprecated(self, used, required):
        """Get the parameter, considering the depreciations."""
        parameter = self.get(required)
        if not parameter:
            if self.get(used):
                parameter = self.get(used)
                print(f'Please use parameter [{required}] instead of',
                      f'[{used}], which is deprecated and will be',
                      f'removed soon from [{self.name}] section',
                      f'of {path} configuration file!')
        return parameter

config = Configurator()


import  numpy as np
import pandas as pd
import logging.handlers
import argparse

LOG_FILENAME = 'data/log/prouct_catalog.log'

# Logger with desired output level
logging.basicConfig(level=1)
log = logging.getLogger("Product_Catalogue")

# Add the log message handler to the logger
handler = logging.handlers.RotatingFileHandler(
    LOG_FILENAME, maxBytes=5000, backupCount=2)

log.addHandler(handler)

#Implement ArgParser
parser = argparse.ArgumentParser(description='Product Catalogue')

parser.add_argument('-i','--input',help='Input file Path')
parser.add_argument('-o','--output',help='Output file Path')
args = parser.parse_args()

log.info("Loading Input Data")
#Load input data into DataFrame
df_articles = pd.read_csv(args.input,header=0,sep="|")

log.info("Pivoting on PART_NO ")
#Pivot on PART_NO
df_articles_pivot = df_articles.pivot(columns="PART_NO",values="COLOR")

df_articles_pivot_frame = df_articles["ARTICLE_CD"].to_frame("ARTICLE_CD")

log.info("Concatenating ARTICLE_CD data with Pivoted PART_NO ")
#Concat ARTICLE_CD data with Pivoted PART_NO
df_result = pd.concat([df_articles_pivot_frame,df_articles_pivot],axis=1)

log.info("Performing Groupby on ARTICLE_CD and max on PART_X")
#Groupby on ARTICLE_CD and max on PART_X
df_groupby = df_result.groupby(by=["ARTICLE_CD"]).max().replace(to_replace=np.NaN,value="")
df_groupby.columns = ['PART_' + str(col) for col in df_groupby.columns if col != "ARTICLE_CD"]

log.info("Writing result to " + args.output)
#Write to csv out
df_groupby.to_csv(args.output)


# -*- coding: utf-8 -*-
import unittest
import os  # noqa: F401
import json  # noqa: F401
import time
import requests

from os import environ
try:
    from ConfigParser import ConfigParser  # py2
except:
    from configparser import ConfigParser  # py3

from pprint import pprint  # noqa: F401

from biokbase.workspace.client import Workspace as workspaceService
from zahmeethContigFilter.zahmeethContigFilterImpl import zahmeethContigFilter
from zahmeethContigFilter.zahmeethContigFilterServer import MethodContext
from zahmeethContigFilter.authclient import KBaseAuth as _KBaseAuth

class zahmeethContigFilterTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        token = environ.get('KB_AUTH_TOKEN', None)
        config_file = environ.get('KB_DEPLOYMENT_CONFIG', None)
        cls.cfg = {}
        config = ConfigParser()
        config.read(config_file)
        for nameval in config.items('zahmeethContigFilter'):
            cls.cfg[nameval[0]] = nameval[1]
        # Getting username from Auth profile for token
        authServiceUrl = cls.cfg['auth-service-url']
        auth_client = _KBaseAuth(authServiceUrl)
        user_id = auth_client.get_user(token)
        # WARNING: don't call any logging methods on the context object,
        # it'll result in a NoneType error
        cls.ctx = MethodContext(None)
        cls.ctx.update({'token': token,
                        'user_id': user_id,
                        'provenance': [
                            {'service': 'zahmeethContigFilter',
                             'method': 'please_never_use_it_in_production',
                             'method_params': []
                             }],
                        'authenticated': 1})
        cls.wsURL = cls.cfg['workspace-url']
        cls.wsClient = workspaceService(cls.wsURL)
        cls.serviceImpl = zahmeethContigFilter(cls.cfg)
        cls.scratch = cls.cfg['scratch']
        cls.callback_url = os.environ['SDK_CALLBACK_URL']

    @classmethod
    def tearDownClass(cls):
        if hasattr(cls, 'wsName'):
            cls.wsClient.delete_workspace({'workspace': cls.wsName})
            print('Test workspace was deleted')

    def getWsClient(self):
        return self.__class__.wsClient

    def getWsName(self):
        if hasattr(self.__class__, 'wsName'):
            return self.__class__.wsName
        suffix = int(time.time() * 1000)
        wsName = "test_zahmeethContigFilter_" + str(suffix)
        ret = self.getWsClient().create_workspace({'workspace': wsName})  # noqa
        self.__class__.wsName = wsName
        return wsName

    def getImpl(self):
        return self.__class__.serviceImpl

    def getContext(self):
        return self.__class__.ctx

    # NOTE: According to Python unittest naming rules test method names should start from 'test'. # noqa

    def test_filter_contigs(self):
        obj_name = "contigset.1"
        contig1 = {'id': '1', 'length': 10, 'md5': 'md5', 'sequence': 'agcttttcat'}
        contig2 = {'id': '2', 'length': 20, 'md5': 'md5', 'sequence': 'agcttttcatagcttttcat'}
        obj = {'contigs': [contig1, contig2], 'id': 'id', 'md5': 'md5', 'name': 'name',
                'source': 'source', 'source_id': 'source_id', 'type': 'type'}
        self.getWsClient().save_objects({'workspace': self.getWsName(), 'objects':
            [{'type': 'KBaseGenomes.ContigSet', 'name': obj_name, 'data': obj}]})
        ret = self.getImpl().filter_contigs(self.getContext(),
                                {
                                    'workspace':self.getWsName(),
                                    'contigset_id':obj_name,
                                    'min_length': 15
                                })

        self.assertEqual(ret[0]['n_initial_contigs'], 2)
        self.assertEqual(ret[0]['n_contigs_removed'], 1)
        self.assertEqual(ret[0]['n_contigs_remaining'], 1)


import enum

import sqlalchemy as sa
from sqlalchemy import orm
from sqlalchemy.dialects import postgresql as psql

from gif_api import db

class GifRating(str, enum.Enum):
    Y = "y"
    G = "g"
    PG = "pg"
    PG13 = "pg-13"
    R = "r"

class Gif(db.Base):
    __tablename__ = "gif"

    gif_id = sa.Column(
        psql.UUID(as_uuid=True),
        primary_key=True,
        unique=True,
        nullable=False,
        server_default=sa.sql.func.gen_random_uuid(),
        doc="This GIF's unique ID.",
    )
    title = sa.Column(
        sa.String(length=127),
        nullable=False,
        doc="The title for this GIF.",
    )
    url = sa.Column(
        sa.String(length=255),
        nullable=False,
        unique=True,
        doc="The unique URL for this GIF.",
    )
    rating = sa.Column(
        sa.Enum(GifRating),
        nullable=False,
        server_default="G",
        doc="The MPAA-style rating for this content.",
    )
    create_datetime = sa.Column(
        sa.DateTime(timezone=True),
        nullable=False,
        server_default=sa.sql.func.now(),
        doc="The date this GIF was added to the database.",
    )
    update_datetime = sa.Column(
        sa.DateTime(timezone=True),
        nullable=False,
        onupdate=sa.sql.func.now(),
        server_default=sa.sql.func.now(),
        doc="The date on which this GIF was last updated.",
    )

    trendings = orm.relationship("Trending", back_populates="trending_gif")

    def __repr__(self) -> str:
        return f"<Gif '{self.gif_id}'>"


import re

from benedict.serializers.abstract import AbstractSerializer
from benedict.serializers.base64 import Base64Serializer
from benedict.serializers.cli import CLISerializer
from benedict.serializers.csv import CSVSerializer
from benedict.serializers.ini import INISerializer
from benedict.serializers.json import JSONSerializer
from benedict.serializers.pickle import PickleSerializer
from benedict.serializers.plist import PListSerializer
from benedict.serializers.query_string import QueryStringSerializer
from benedict.serializers.toml import TOMLSerializer
from benedict.serializers.xls import XLSSerializer
from benedict.serializers.xml import XMLSerializer
from benedict.serializers.yaml import YAMLSerializer

__all__ = [
    "AbstractSerializer",
    "Base64Serializer",
    "CLISerializer",
    "CSVSerializer",
    "INISerializer",
    "JSONSerializer",
    "PickleSerializer",
    "PListSerializer",
    "QueryStringSerializer",
    "TOMLSerializer",
    "XLSSerializer",
    "XMLSerializer",
    "YAMLSerializer",
]

_BASE64_SERIALIZER = Base64Serializer()
_CLI_SERIALIZER = CLISerializer()
_CSV_SERIALIZER = CSVSerializer()
_INI_SERIALIZER = INISerializer()
_JSON_SERIALIZER = JSONSerializer()
_PICKLE_SERIALIZER = PickleSerializer()
_PLIST_SERIALIZER = PListSerializer()
_QUERY_STRING_SERIALIZER = QueryStringSerializer()
_TOML_SERIALIZER = TOMLSerializer()
_XLS_SERIALIZER = XLSSerializer()
_XML_SERIALIZER = XMLSerializer()
_YAML_SERIALIZER = YAMLSerializer()

_SERIALIZERS_LIST = [
    _BASE64_SERIALIZER,
    _CLI_SERIALIZER,
    _CSV_SERIALIZER,
    _INI_SERIALIZER,
    _JSON_SERIALIZER,
    _PICKLE_SERIALIZER,
    _PLIST_SERIALIZER,
    _QUERY_STRING_SERIALIZER,
    _TOML_SERIALIZER,
    _YAML_SERIALIZER,
    _XLS_SERIALIZER,
    _XML_SERIALIZER,
]

_SERIALIZERS_BY_EXTENSION = {}
for serializer in _SERIALIZERS_LIST:
    for extension in serializer.extensions():
        _SERIALIZERS_BY_EXTENSION[extension] = serializer

_SERIALIZERS_EXTENSIONS = [
    f".{extension}" for extension in _SERIALIZERS_BY_EXTENSION.keys()
]

def get_format_by_path(path):
    path = str(path)
    path = path.lower()
    for extension in _SERIALIZERS_EXTENSIONS:
        if path.endswith(extension):
            return extension[1:]
    return None

def get_serializer_by_format(format):
    format_key = (format or "").lower().strip()
    format_key = re.sub(r"[\s\-\_]*", "", format_key)
    serializer = _SERIALIZERS_BY_EXTENSION.get(format_key, None)
    return serializer


"""Implementation of prototypical networks for Omniglot."""

import argparse
import os

import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
from torch.utils import tensorboard

import omniglot
# import util

NUM_INPUT_CHANNELS = 1
NUM_HIDDEN_CHANNELS = 32
KERNEL_SIZE = 3
NUM_CONV_LAYERS = 4
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
SUMMARY_INTERVAL = 10
SAVE_INTERVAL = 100
PRINT_INTERVAL = 10
VAL_INTERVAL = PRINT_INTERVAL * 5
NUM_TEST_TASKS = 600

class ProtoNetNetwork(nn.Module):
    """Container for ProtoNet weights and image-to-latent computation."""

    def __init__(self):
        """Inits ProtoNetNetwork.

        The network consists of four convolutional blocks, each comprising a
        convolution layer, a batch normalization layer, ReLU activation, and 2x2
        max pooling for downsampling. There is an additional flattening
        operation at the end.

        Note that unlike conventional use, batch normalization is always done
        with batch statistics, regardless of whether we are training or
        evaluating. This technically makes meta-learning transductive, as
        opposed to inductive.
        """
        super().__init__()

        in_channels = NUM_INPUT_CHANNELS

        # for _ in range(NUM_CONV_LAYERS):
        #     layers.append(
        #         nn.Conv2d(
        #             in_channels,
        #             NUM_HIDDEN_CHANNELS,
        #             (KERNEL_SIZE, KERNEL_SIZE),
        #             padding='same'
        #         )
        #     )
        #     layers.append(nn.BatchNorm2d(NUM_HIDDEN_CHANNELS))
        #     layers.append(nn.ReLU())
        #     layers.append(nn.MaxPool2d(2))
        #     in_channels = NUM_HIDDEN_CHANNELS

        self.conv1 = nn.Conv2d(
            in_channels,
            NUM_HIDDEN_CHANNELS,
            (KERNEL_SIZE, KERNEL_SIZE),
            padding='same'
        )
        torch.nn.init.normal_(self.conv1.weight)
        torch.nn.init.normal_(self.conv1.bias)

        # self._layers = nn.Sequential(self.conv1, nn.ReLU(), nn.AvgPool2d(), nn.Flatten())

        self.to(DEVICE)

    # def forward(self, images):
    #     """Computes the latent representation of a batch of images.

    #     Args:
    #         images (Tensor): batch of Omniglot images
    #             shape (num_images, channels, height, width)

    #     Returns:
    #         a Tensor containing a batch of latent representations
    #             shape (num_images, latents)
    #     """
    #     return self._layers(images)

    def forward_(self, images, weight, bias):

        out = torch.nn.functional.conv2d(images, weight, bias, padding='same')
        out = torch.nn.functional.relu(out)
        out = torch.nn.functional.avg_pool2d(out, 2)

        return out

    def constr_(self, images, weight, bias):

        out = torch.nn.functional.conv2d(images, weight, bias, padding='same')
        out = 2*torch.nn.functional.relu(out) - out

        return out

class ProtoNet:
    """Trains and assesses a prototypical network."""

    def __init__(self, learning_rate, log_dir):
        """Inits ProtoNet.

        Args:
            learning_rate (float): learning rate for the Adam optimizer
            log_dir (str): path to logging directory
        """

        self._network = ProtoNetNetwork()
        self._optimizer = torch.optim.Adam(
            self._network.parameters(),
            lr=learning_rate
        )
        self._log_dir = log_dir
        os.makedirs(self._log_dir, exist_ok=True)

        self._start_train_step = 0

        self.constraints_all = []
        self.distances_query = []
        self.labels_query = []

    def _step(self, task_batch):
        """Computes ProtoNet mean loss (and accuracy) on a batch of tasks.

        Args:
            task_batch (tuple[Tensor, Tensor, Tensor, Tensor]):
                batch of tasks from an Omniglot DataLoader

        Returns:
            a Tensor containing mean ProtoNet loss over the batch
                shape ()
            mean support set accuracy over the batch as a float
            mean query set accuracy over the batch as a float
        """
        loss_batch = []
        accuracy_support_batch = []
        accuracy_query_batch = []
        for i, task in enumerate(task_batch):

            print(f'task: {i}')

            images_support, labels_support, images_query, labels_query = task
            images_support = images_support.to(DEVICE)
            labels_support = labels_support.to(DEVICE)
            images_query = images_query.to(DEVICE)
            labels_query = labels_query.to(DEVICE)

            # ********************************************************
            # ******************* YOUR CODE HERE *********************
            # ********************************************************
            # TODO: finish implementing this method.
            # For a given task, compute the prototypes and the protonet loss.
            # Use F.cross_entropy to compute classification losses.
            # Use util.score to compute accuracies.
            # Make sure to populate loss_batch, accuracy_support_batch, and
            # accuracy_query_batch.

            K = (labels_support==0).sum().item() # K-shot (support)
            K_query = (labels_query==0).sum().item() # K-shot (query)

            N = torch.max(labels_support).item()+1 # N-way

            # embeddings_support = self._network.forward(images_support) # (N * K, latents)
            # embeddings_query = self._network.forward(images_query) # (N * K_query, latents)

            # embeddings_support_jac = torch.autograd.functional.jacobian(lambda w, b : self._network.forward_(images_support, w, b), (self._network.conv1.weight, self._network.conv1.bias))
            # embeddings_query_jac = torch.autograd.functional.jacobian(lambda w, b : self._network.forward_(images_query, w, b), (self._network.conv1.weight, self._network.conv1.bias))

            # print("constraints!")

            # support_constr = torch.autograd.functional.jacobian(
            #     lambda w, b : self._network.constr_(images_support, w, b),
            #     (self._network.conv1.weight, self._network.conv1.bias),
            #     strategy='forward-mode'
            # )

            all_images = torch.concat((images_support, images_query))
            all_masks = torch.nn.functional.conv2d(all_images, self._network.conv1.weight, self._network.conv1.bias)

            # images_query_masks = torch.nn.functional.conv2d(images_query, self._network.conv1.weight, self._network.conv1.bias)

            # (2*images_support_masks - 1)

            # support_conv = [torch.autograd.functional._jacfwd(
            #     lambda w, b : torch.nn.functional.conv2d(all_images, w, b),
            #     (self._network.conv1.weight[j:j+1], self._network.conv1.bias[j:j+1]),
            #     vectorize=True
            # ) for j in range(NUM_HIDDEN_CHANNELS)]

            all_conv = []

            # lazy generation of of conv2d jacobian wrt weights

            for j in range(NUM_HIDDEN_CHANNELS):

                grad_w, grad_b = torch.autograd.functional._jacfwd(
                    lambda w, b : torch.nn.functional.conv2d(all_images, w, b),
                    (self._network.conv1.weight[j:j+1], self._network.conv1.bias[j:j+1]),
                    vectorize=True
                )

                grad_w = torch.flatten(grad_w, start_dim=4)
                grad_b = torch.flatten(grad_b, start_dim=4)

                all_conv.append(torch.concat((grad_w, grad_b), dim=4))

            all_conv = torch.stack(all_conv, dim=1).squeeze()

            constraint = ((2*all_masks-1).unsqueeze(-1)*all_conv) # num_images * hidden_layers * h * w * weights_per_kernel

            num_images, hidden_layers, h, w, weights_per_kernel = constraint.shape

            self.constraints_all.append(constraint.detach().cpu())

            all_after_relu = all_masks.unsqueeze(-1)*all_conv # num_images * hidden_layers * h * w * weights_per_kernel

            all_before_pooling = all_after_relu.permute(0, 1, 4, 2, 3).reshape(num_images, hidden_layers*weights_per_kernel, h, w)

            all_after_pooling = torch.nn.functional.avg_pool2d(all_before_pooling, 2)

            all_after_pooling = all_after_pooling.permute(0, 2, 3, 1).reshape(num_images, -1, hidden_layers, weights_per_kernel)

            # print(all_after_pooling.shape)

            embeddings_support = all_after_pooling[:N*K]
            embeddings_query = all_after_pooling[N*K:]

            prototypes = torch.mean(embeddings_support.reshape(N, K, *embeddings_support.shape[1:]), axis=1) # (N, latents)

            distances_query = embeddings_query[:, None, ...] - prototypes

            self.distances_query.append(distances_query.detach().cpu())
            self.labels_query.append(labels_query.detach().cpu())

            # print(pos_norm.shape)
            # print(distances_query.shape)

            # exit()

            # print(images_support_masks.shape)
            # print(support_conv[0][0].shape)
            # exit()

            # support_constr = [torch.autograd.functional._jacfwd(
            #     lambda w, b : self._network.constr_(images_support, w, b),
            #     (self._network.conv1.weight[j:j+1], self._network.conv1.bias[j:j+1]),
            #     vectorize=True
            # ) for j in range(NUM_HIDDEN_CHANNELS)]

            # query_constr = [torch.autograd.functional._jacfwd(
            #     lambda w, b : self._network.constr_(images_query, w, b),
            #     (self._network.conv1.weight[j:j+1], self._network.conv1.bias[j:j+1]),
            #     vectorize=True
            # ) for j in range(NUM_HIDDEN_CHANNELS)]

            # def func(weight, bias):

            #     embeddings_support = torch.flatten(self._network.forward_(images_support, weight, bias), start_dim=1)
            #     print(embeddings_support.shape)
            #     embeddings_query = torch.flatten(self._network.forward_(images_query, weight, bias), start_dim=1)

            #     prototypes = torch.mean(embeddings_support.reshape((N, K, -1)), axis=1) # (N, latents)

            #     distances_support = (embeddings_support[:, None, :] - prototypes)
            #     distances_query = (embeddings_query[:, None, :] - prototypes)

            #     return distances_support, distances_query

            # print("norms!")

            # pos_norm, neg_norm = [torch.autograd.functional.jacobian(
            #     func,
            #     (self._network.conv1.weight[j:j+1], self._network.conv1.bias[j:j+1])
            # ) for j in range(NUM_HIDDEN_CHANNELS)]

            # print(pos_norm.shape)
            # print(neg_norm.shape)

            # self.constraints.append(support_constr)
            # self.constraints.append(query_constr)
            # self.pos_norms.append(pos_norm)
            # self.neg_norms.append(neg_norm)

    def train(self, dataloader_train, dataloader_val, writer):
        """Train the ProtoNet.

        Consumes dataloader_train to optimize weights of ProtoNetNetwork
        while periodically validating on dataloader_val, logging metrics, and
        saving checkpoints.

        Args:
            dataloader_train (DataLoader): loader for train tasks
            dataloader_val (DataLoader): loader for validation tasks
            writer (SummaryWriter): TensorBoard logger
        """
        print(f'Starting training at iteration {self._start_train_step}.')
        for i_step, task_batch in enumerate(
                dataloader_train,
                start=self._start_train_step
        ):
            # self._optimizer.zero_grad()
            # loss, accuracy_support, accuracy_query = self._step(task_batch)
            # loss.backward()
            # self._optimizer.step()

            self._step(task_batch)

            if i_step % PRINT_INTERVAL == 0:
                print(
                    f'Iteration {i_step}: '
                    # f'loss: {loss.item():.3f}, '
                    # f'support accuracy: {accuracy_support.item():.3f}, '
                    # f'query accuracy: {accuracy_query.item():.3f}'
                )
                # writer.add_scalar('loss/train', loss.item(), i_step)
                # writer.add_scalar(
                #     'train_accuracy/support',
                #     accuracy_support.item(),
                #     i_step
                # )
                # writer.add_scalar(
                #     'train_accuracy/query',
                #     accuracy_query.item(),
                #     i_step
                # )

            # if i_step % VAL_INTERVAL == 0:
            #     with torch.no_grad():
            #         losses, accuracies_support, accuracies_query = [], [], []
            #         for val_task_batch in dataloader_val:
            #             loss, accuracy_support, accuracy_query = (
            #                 self._step(val_task_batch)
            #             )
            #             losses.append(loss.item())
            #             accuracies_support.append(accuracy_support)
            #             accuracies_query.append(accuracy_query)
            #         loss = np.mean(losses)
            #         accuracy_support = np.mean(accuracies_support)
            #         accuracy_query = np.mean(accuracies_query)
            #     print(
            #         f'Validation: '
            #         f'loss: {loss:.3f}, '
            #         f'support accuracy: {accuracy_support:.3f}, '
            #         f'query accuracy: {accuracy_query:.3f}'
            #     )
            #     writer.add_scalar('loss/val', loss, i_step)
            #     writer.add_scalar(
            #         'val_accuracy/support',
            #         accuracy_support,
            #         i_step
            #     )
            #     writer.add_scalar(
            #         'val_accuracy/query',
            #         accuracy_query,
            #         i_step
            #     )

            # if i_step % SAVE_INTERVAL == 0:
            #     self._save(i_step)

    def test(self, dataloader_test):
        """Evaluate the ProtoNet on test tasks.

        Args:
            dataloader_test (DataLoader): loader for test tasks
        """
        accuracies = []
        for task_batch in dataloader_test:
            accuracies.append(self._step(task_batch)[2])
        mean = np.mean(accuracies)
        std = np.std(accuracies)
        mean_95_confidence_interval = 1.96 * std / np.sqrt(NUM_TEST_TASKS)
        print(
            f'Accuracy over {NUM_TEST_TASKS} test tasks: '
            f'mean {mean:.3f}, '
            f'95% confidence interval {mean_95_confidence_interval:.3f}'
        )

    def load(self, checkpoint_step):
        """Loads a checkpoint.

        Args:
            checkpoint_step (int): iteration of checkpoint to load

        Raises:
            ValueError: if checkpoint for checkpoint_step is not found
        """
        target_path = (
            f'{os.path.join(self._log_dir, "state")}'
            f'{checkpoint_step}.pt'
        )
        if os.path.isfile(target_path):
            state = torch.load(target_path)
            self._network.load_state_dict(state['network_state_dict'])
            self._optimizer.load_state_dict(state['optimizer_state_dict'])
            self._start_train_step = checkpoint_step + 1
            print(f'Loaded checkpoint iteration {checkpoint_step}.')
        else:
            raise ValueError(
                f'No checkpoint for iteration {checkpoint_step} found.'
            )

    def _save(self, checkpoint_step):
        """Saves network and optimizer state_dicts as a checkpoint.

        Args:
            checkpoint_step (int): iteration to label checkpoint with
        """
        torch.save(
            dict(network_state_dict=self._network.state_dict(),
                 optimizer_state_dict=self._optimizer.state_dict()),
            f'{os.path.join(self._log_dir, "state")}{checkpoint_step}.pt'
        )
        print('Saved checkpoint.')

def main(args):
    log_dir = args.log_dir
    if log_dir is None:
        log_dir = f'./logs/protonet/omniglot.way:{args.num_way}.support:{args.num_support}.query:{args.num_query}.lr:{args.learning_rate}.batch_size:{args.batch_size}'  # pylint: disable=line-too-long
    print(f'log_dir: {log_dir}')
    # writer = tensorboard.SummaryWriter(log_dir=log_dir)
    writer = None

    protonet = ProtoNet(args.learning_rate, log_dir)

    if args.checkpoint_step > -1:
        protonet.load(args.checkpoint_step)
    else:
        print('Checkpoint loading skipped.')

    if not args.test:
        num_training_tasks = args.batch_size * (args.num_train_iterations -
                                                args.checkpoint_step - 1)
        print(
            f'Training on tasks with composition '
            f'num_way={args.num_way}, '
            f'num_support={args.num_support}, '
            f'num_query={args.num_query}'
        )
        dataloader_train = omniglot.get_omniglot_dataloader(
            'train',
            args.batch_size,
            args.num_way,
            args.num_support,
            args.num_query,
            num_training_tasks
        )
        dataloader_val = omniglot.get_omniglot_dataloader(
            'val',
            args.batch_size,
            args.num_way,
            args.num_support,
            args.num_query,
            args.batch_size * 4
        )
        protonet.train(
            dataloader_train,
            dataloader_val,
            writer
        )
        labels_query = torch.stack(protonet.labels_query).numpy()
        constraints_all = torch.stack(protonet.constraints_all).numpy()
        distances_query = torch.stack(protonet.distances_query).numpy()
        weight = protonet._network.conv1.weight.detach().cpu().numpy()
        bias = protonet._network.conv1.bias.detach().cpu().numpy()

        print('saving to disk!')
        np.savez(
            'protonet2.npz',
            labels_query=labels_query,
            constraints_all=constraints_all,
            distances_query=distances_query,
            weight=weight,
            bias=bias
        )

    else:
        print(
            f'Testing on tasks with composition '
            f'num_way={args.num_way}, '
            f'num_support={args.num_support}, '
            f'num_query={args.num_query}'
        )
        dataloader_test = omniglot.get_omniglot_dataloader(
            'test',
            1,
            args.num_way,
            args.num_support,
            args.num_query,
            NUM_TEST_TASKS
        )
        protonet.test(dataloader_test)

if __name__ == '__main__':
    parser = argparse.ArgumentParser('Train a ProtoNet!')
    parser.add_argument('--log_dir', type=str, default=None,
                        help='directory to save to or load from')
    parser.add_argument('--num_way', type=int, default=5,
                        help='number of classes in a task')
    parser.add_argument('--num_support', type=int, default=1,
                        help='number of support examples per class in a task')
    parser.add_argument('--num_query', type=int, default=15,
                        help='number of query examples per class in a task')
    parser.add_argument('--learning_rate', type=float, default=0.001,
                        help='learning rate for the network')
    parser.add_argument('--batch_size', type=int, default=16,
                        help='number of tasks per outer-loop update')
    parser.add_argument('--num_train_iterations', type=int, default=1,
                        help='number of outer-loop updates to train for')
    parser.add_argument('--test', default=False, action='store_true',
                        help='train or test')
    parser.add_argument('--checkpoint_step', type=int, default=-1,
                        help=('checkpoint iteration to load for resuming '
                              'training, or for evaluation (-1 is ignored)'))

    main_args = parser.parse_args()
    main(main_args)


#Assignment 6
#Problem 2
#Siddharth Dixit
#2018ME20727

import numpy as np
import math
import matplotlib.pyplot as plt
import scipy

from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_absolute_error, mean_squared_error
from sklearn.metrics import r2_score
from sklearn.metrics import explained_variance_score

from scipy.stats import norm

n = 10000

# a = 0.2 & b = 0.3
# ab =[a, b]
ab = [0.2, 0.3]

#mean 0 and standard deviation 0.5
#msd = [mean, standard deviation]
msd = [0, 0.5]

X = np.linspace(1,5000,n)
s = np.random.normal(msd[0], msd[1], n)
y = []

x = 0
while x < n:
    y.append(ab[0]+ab[1]*X[x]+s[x])
    x = x + 1

x_train,x_test,y_train,y_test=train_test_split(X.reshape(-1, 1),y,test_size=0.2)

regression = LinearRegression().fit(x_train, y_train)
print("Regression Score =",regression.score(x_test,y_test))
intercept = regression.intercept_
slope = regression.coef_

y_pred = []
for xi in x_test:
    y_pred.append(regression.predict(xi.reshape(-1, 1))[0])

residual_error = []

i = 0
while i < len(y_test):
    residual_error.append(y_test[i] - y_pred[i])
    i = i + 1

X_axis = np.arange(len(x_test))
plt.bar(X_axis, residual_error)
plt.xlabel("X values")
plt.ylabel("Residual error values")
plt.title("Residual error histogram")
plt.cla()
mean, std = norm.fit(residual_error)
print("Mean =",mean)
print("Standard Deviation =", std)
plt.hist(residual_error, bins=25, density=True, alpha=0.6, color='orange')
xmin, xmax = plt.xlim()
x = np.linspace(xmin, xmax, 100)
p = norm.pdf(x, mean, std)
plt.plot(x, p, 'k', linewidth=2)
title = "2018ME20727 : Mean = %.2f,  standard deviation = %.2f" % (mean, std)
plt.title(title)

plt.show()


import aiofiles
import os
import pyfiglet
from rich import print
from time import perf_counter
import pyfiglet
from util.const import *

def cls():
    os.system('cls' if os.name == 'nt' else 'clear')

async def load_tokens(file_path, token_queue):
    line_count = 0  # Initialize a counter for the lines
    async with aiofiles.open(file_path, mode='r') as f:
        async for line in f:
            line = line.strip()
            await token_queue.put(line)
            line_count += 1  # Increment the counter for each line
    return line_count  # Return the total number of lines read

def banner():
    global total_tokens

    os.system("title FatBee's Account Checker")
    cls()
    ascii_banner = pyfiglet.figlet_format("FatBee's Checker")
    print(f"[yellow]{ascii_banner}")
    print(
        f"[bold white]🐝 Made by [bold yellow]FatBee[/bold yellow]  |  💬 Telegram: [bold cyan]@fatbeebhw[/bold cyan]  |  💬 Telegram Group: [bold cyan]@twitteropensource[/bold cyan]  | ✅ Version: [bold light_green]{VERSION}[/bold light_green]"
    )

    try:
        with open("tokens.txt", "r") as f:
            tokens = f.read().splitlines()
        total_tokens = len(tokens)

        if not tokens:
            print("[yellow]\n[!] Tokens File is Empty.\n")
            exit()
        else:
            print(
                f"[yellow]\n[!] Loaded [cyan]{total_tokens} [yellow]Accounts")

        # Data Usage
        initial_data_usage = total_tokens * 1.1
        expected_data_usage = round(initial_data_usage / 1000, 3)
        print(
            f"[yellow][!] Expected Data Usage is less than [cyan]{expected_data_usage} MB")

        # Expected Time to Complete
        initial_time = total_tokens * 0.035
        minutes, seconds = divmod(initial_time, 60)
        rounded_minutes = round(minutes)
        rounded_seconds = round(seconds)
        print(
            f"[yellow][!] Expected Run Time: [cyan]{rounded_minutes} Minute(s) {rounded_seconds} Second(s)")

        if PROXY_URL:
            print("[bold green][*] You are checking with proxy.")
        else:
            print("[bold red][!] You are checking proxyless!.")

        print(
            "[yellow]━━━━━━━━━━━━━━━━ [bold cyan]Initializing Threads [yellow]━━━━━━━━━━━━━━━━ \n")

    except FileNotFoundError:
        print("[red]\n[!] Tokens.txt file not found.")
        exit()

    return total_tokens  # Return the total number of tokens

def check_completed(t1_start, total_tokens, total_valid, total_dead, total_locked, total_consent,total_suspended):
    cls()
    counts = {
        "VALID": total_valid,
        "DEAD": total_dead,
        "LOCKED": total_locked,
        "CONSENT": total_consent,
        "SUSPENDED": total_suspended
    }

    total_bad = total_dead + total_locked + total_suspended + total_consent
    if total_tokens == 0:
        percent = 0  # Avoid dividing by zero
    else:
        percent = (total_bad / total_tokens) * 100
    percent = round(percent, 3)
    final_time = round(perf_counter() - t1_start)
    mm, ss = divmod(final_time, 60)

    ascii_banner = pyfiglet.figlet_format("FatBee's Checker")
    print(f"[bold yellow]{ascii_banner}")
    print(
        f"[bold white]🐝 Made by [bold yellow]FatBee[/bold yellow]  |  💬 Telegram: [bold cyan]@fatbeebhw[/bold cyan]  |  💬 Telegram Group: [bold cyan]@twitteropensource[/bold cyan]  | ✅ Version: [bold light_green]{VERSION}[/bold light_green]")
    print(f"[bold white]━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[/bold white]")
    print(
        f"[bold green]⏳ Total Running Time: {mm} Minute(s) {ss} Second(s)[/bold green]")
    print(
        f"[bold yellow][*] [bold white]Accounts Checked: [bold yellow]{total_tokens}[/bold yellow]")
    print(
        f"[bold yellow][*] [bold white]Lock Rate: [bold red]{percent}%[/bold red]")
    print(
        f"[bold yellow][*] [bold white]Valid Accounts: [bold green]{counts['VALID']}[/bold green]")
    print(
        f"[bold yellow][*] [bold white]Invalid Accounts: [bold red]{counts['DEAD']}[/bold red]")
    print(
        f"[bold yellow][*] [bold white]Suspended Accounts: [bold red]{counts['SUSPENDED']}[/bold red]")
    print(
        f"[bold yellow][*] [bold white]Unlockable: [bold yellow]{counts['LOCKED']}[/bold yellow]")
    print(
        f"[yellow][*] [bold white]Consent Locked: [bold yellow]{counts['CONSENT']}[/bold yellow]")
    print(f"[bold white]━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[/bold white]")
    print("[bold red] ** RESTARTING THE CHECKER WILL DELETE ALL THE FILES IN OUTPUT FOLDER **[/bold red]")
    input("Press any key to exit..")

async def cleanup_files(file_names):
    for file_name in file_names:
        async with aiofiles.open(f"output/{file_name}.txt", 'w') as f:
            await f.write('')


import os
import pathlib

import uvicorn
from dotenv import load_dotenv

from fastapi import FastAPI
from fastapi_sqlalchemy import DBSessionMiddleware, db

from src.models.library_models import Book, Author
from src.schemas.library_post_schemas import BookSchema, AuthorSchema

# Define environment variables for customization of alembic
BASE_DIR = pathlib.Path(__file__).absolute().parent.parent
load_dotenv(BASE_DIR / ".env")

# Create FastAPI instance
app = FastAPI()
app.add_middleware(DBSessionMiddleware, db_url=os.environ["DATABASE_URL"])

# Define api endpoints
@app.get("/")
async def root():
    return {
        "message": "Hello to my fastapi world! Nothing to see here, try using /docs"
    }

@app.post("/add-book/", response_model=BookSchema)
def add_book(book: BookSchema):
    db_book = Book(
        title=book.title, rating=book.rating, author_id=book.author_id
    )
    db.session.add(db_book)
    db.session.commit()
    return db_book

@app.post("/add-author/", response_model=AuthorSchema)
def add_author(author: AuthorSchema):
    db_author = Author(name=author.name, age=author.age)
    db.session.add(db_author)
    db.session.commit()
    return db_author

@app.get("/books/")
def get_books():
    books = db.session.query(Book).all()
    return books

if (
    __name__ == "__main__"
):  # this is to be able to run the api on local machine, no docker containers
    uvicorn.run(app, host="0.0.0.0", port=8000)