# -*- coding: utf-8 -*-
#!/usr/bin/env python
"""A class for weighted wavelet z-transform analysis.
"""

from datetime import datetime
import numpy as np
import pickle
from scipy.interpolate import interp1d, splrep, splev
import sys

__author__ = "Sebastian Kiehlmann"
__credits__ = ["Sebastian Kiehlmann", "Walther Max-Moerbeck", "Oliver King"]
__license__ = "BSD"
__maintainer__ = "Sebastian Kiehlmann"
__email__ = "skiehlmann@mail.de"
__status__ = "Production"

#==============================================================================
# CLASSES
#==============================================================================

class WWZ:
    """A class for weighted wavelet z-transform analysis."""

    #--------------------------------------------------------------------------
    def __init__(self, time=None, flux=None, flux_err=0):
        """A class for weighted wavelet z-transform analysis.

        Parameters
        ----------
        time : array-like, optional
            Time stamps of the input data. The default is None.
        flux : array-like, optional
            Flux of the input data. The default is None.
        flux_err : array-like or float, optional
            Flux uncertainty of the input data. The default is 0.

        Returns
        -------
        None.

        Notes
        -----
        Setting the data during the instance initialization is optional. Data
        can be provided with the set_data() method.
        """

        print('WWZ instance created.')

        if time is not None and flux is not None:
            self.set_data(time, flux, flux_err=flux_err)
        else:
            self._reset()

    #--------------------------------------------------------------------------
    def _reset(self):
        """Reset all class attributes to None.

        Returns
        -------
        None.
        """

        # data time, flux, and flux error:
        self.t = None
        self.x = None
        self.s_x = None

        # number of data points:
        self.n = 0

        # attributes filled by set_freq() and set_tau():
        self.linear_period = False
        self.freq = None
        self.tau = None
        self.len0 = None
        self.len1 = None

        # class attributes filled by transform method:
        self.snr_weights = None
        self.c = None
        self.v_x = None
        self.v_y = None
        self.n_eff = None
        self.wwz = None
        self.wwa = None
        self.wwz_sig = None
        self.wwa_sig = None
        self.wwz_pval = None
        self.wwa_pval = None
        self.n_sig = None

    #--------------------------------------------------------------------------
    def set_data(self, time, flux, flux_err=0.0, verbose=True):
        """Provide data for the analysis.

        Parameters
        ----------
        time : array-like, optional
            Time stamps of the input data. The default is None.
        flux : array-like, optional
            Flux of the input data. The default is None.
        flux_err : array-like or float, optional
            Flux uncertainty of the input data. The default is 0.
        verbose : bool, optional
            Turns printing of information on or off. The default is True.

        Returns
        -------
        None.

        Notes
        -----
        When new data is provided all class attributes that may contain the
        results of an earlier analysis will be reset to None.
        """

        self._reset()

        # data time, flux, and flux error:
        self.t = np.asarray(time)
        self.x = np.asarray(flux)
        self.s_x = np.asarray(flux_err)

        # number of data points:
        self.n = len(time)

        if verbose:
            print('Data stored.')

    #--------------------------------------------------------------------------
    def get_period_lims(self, p_min_factor=4., p_max_factor=5.):
        """Get suggested limits for the periods.

        Parameters
        ----------
        p_min_factor : float, optional
            If p_min is None, p_min is calculated as the median difference
            between time steps multiplied with this factor.
        p_max_factor : float, optional
            If p_max is None, p_max is calculated as the total time of the data
            devided by this factor.

        Returns
        -------
        p_min : float
            Suggested minimum period.
        p_max : float
            Suggested maximum period.
        """

        p_min = np.round(np.median(np.diff(self.t)) * p_min_factor)
        p_max = np.round((self.t[-1] - self.t[0]) / p_max_factor)

        return p_min, p_max

    #--------------------------------------------------------------------------
    def get_freq(
            self, p_min=None, p_max=None, n_bins=100, diff=None,
            linear_period=False, p_min_factor=4., p_max_factor=5.):
        """Create a range of frequencies for the analysis.

        Parameters
        ----------
        p_min : float, optional
            Minimum period. If provided, this shorest period will define the
            highest frequency. If None, the shortest period will be determined
            by the median of the time difference between data points. The
            default is None.
        p_max : float, optional
            Maximum period. If provided, this largest period will define the
            lowest frequency. If None, the largest period will be determined
            by the total duration of the data. The default is None.
        n_bins : int, optional
            Number of frequencies. The default is 100.
        diff : float, optional
            Sets the difference between frequency bins. The number of bins will
            be calculated accordingly. If linear_period=True, this sets the
            difference between periods. If diff is set, this overwrites any
            input to n_bins. The default is None.
        linear_period : bool, optional
            If False, the frequencies will be distributed in equal steps. If
            True, the frequencies are chosen such that the corresponding
            periods are distributed in equal steps. The default is False.
        p_min_factor : float, optional
            If p_min is None, p_min is calculated as the median difference
            between time steps multiplied with this factor.
        p_max_factor : float, optional
            If p_max is None, p_max is calculated as the total time of the data
            devided by this factor.

        Raises
        ------
        ValueError
            Error is raise when neither n_bins nor diff is provided.

        Returns
        -------
        freq : numpy.ndarray
            One dimensional array of frequencies.
        """

        p_min_suggested, p_max_suggested =  self.get_period_lims(
                p_min_factor=p_min_factor, p_max_factor=p_max_factor)

        if p_min is None:
            p_min = p_min_suggested
        elif p_min < p_min_suggested:
            print('WARNING: p_min is shorter than {0} '.format(p_min_factor),
                  'times the median time step of the data.\np_min should ',
                  'not be shorter than {0}.\n'.format(p_min_suggested))

        if p_max is None:
            p_max = p_max_suggested
        elif p_max > p_max_suggested:
            print('WARNING: p_max is larger than the maximum time range',
                  'of the data devided by {0}.\np_max '.format(p_max_factor),
                  'should not be larger than {0}.\n'.format(p_max_suggested))

        freq_min = 1. / p_max
        freq_max = 1. / p_min

        if linear_period:
            if diff is not None:
                period = np.arange(p_min, p_max, diff)

            elif isinstance(n_bins, int):
                period = np.linspace(p_min, p_max, n_bins)
                diff = period[1] - period[0]

            else:
                raise ValueError("Either n_bins or diff needs to be set.")

            freq = 1. / period
            freq = freq[::-1]

        else:
            if diff is not None:
                freq = np.arange(freq_min, freq_max, diff)

            elif isinstance(n_bins, int):
                freq = np.linspace(freq_min, freq_max, n_bins)
                diff = freq[1] - freq[0]

            else:
                raise ValueError("Either n_bins or diff needs to be set.")

        print('Linear range of frequencies created with')
        print('Shortest period:       {0:10.2e}'.format(p_min))
        print('Longest period:        {0:10.2e}'.format(p_max))
        if linear_period:
            print('Period interval:       {0:10.2e}'.format(diff))
        else:
            print('Period interval:       non-linear')
        print('Lowest frequency:      {0:10.2e}'.format(freq_min))
        print('Highest frequency:     {0:10.2e}'.format(freq_max))
        if linear_period:
            print('Frequency interval:    non-linear')
        else:
            print('Frequency interval:    {0:10.2e}'.format(diff))
        print('Number of frequencies: {0:10d}\n'.format(freq.size))

        return freq

    #--------------------------------------------------------------------------
    def get_tau(self, t_min=None, t_max=None, n_div=8, n_bins=None, dtau=None):
        """Create a range of taus (time points) for the analysis.

        Parameters
        ----------
        t_min : float, optional
            Earliest time point. If not provided, will be the first data point.
            The default is None.
        t_max : float, optional
            Latest time point. If not provided, will be the last data point.
            The default is None.
        n_div : int, optional
            The total time of the data will be devided by n_div. The result is
            converted to the nearest integer. This sets the number of time
            bins. A larger n_div will result in fewer time bins, and vice
            versa. The default is 8.
        n_bins : int, optional
            Number of time bins. If n_bins is set, this overwrites any input to
            n_div. The default is None.
        dtau : float, optional
            Sets the difference between time bins. The number of bins will
            be calculated accordingly. If diff is set, this overwrites any
            input to n_bins and n_div. The default is None.

        Raises
        ------
        ValueError
            Error is raise when neither n_bins nor diff is provided.

        Returns
        -------
        freq : numpy.ndarray
            One dimensional array of frequencies.
        """

        t_min = self.t[0]
        t_max = self.t[-1]

        if dtau is not None:
            tau = np.arange(t_min, t_max, dtau)
            n_bins = tau.size

        elif isinstance(n_bins, int):
            tau = np.linspace(t_min, t_max, n_bins)
            dtau = tau[1] - tau[0]

        elif isinstance(n_div, int):
            n_bins = int((t_max - t_min) / n_div)
            tau = np.linspace(t_min, t_max, n_bins)
            dtau = tau[1] - tau[0]

        else:
            raise ValueError("Either n_div, n_bins, or dtau needs to be set.")

        print('Linear range of tau (time points) created with')
        print('Earliest time:         {0:10.1f}'.format(t_min))
        print('Latest time:           {0:10.1f}'.format(t_max))
        print('Time interval:         {0:10.1f}'.format(dtau))
        print('Points in time:        {0:10d}\n'.format(n_bins))

        return tau

    #--------------------------------------------------------------------------
    def set_freq(self, freq, verbose=True):
        """Set the frequencies for the analysis.

        Parameters
        ----------
        freq : array-like
            A range of frequencies at which the WWZ transform will be
            calculated.
        verbose : bool, optional
            Turns printing of information on or off. The default is True.

        Returns
        -------
        None.

        Notes
        -----
        It is recommended to use the get_freq() method to create the array of
        frequencies.
        """

        self.freq = np.asarray(freq)
        self.len1 = freq.size

        # check if frequencies are linear in period space:
        dperiod = np.diff(1. / freq)
        self.linear_period = np.all(np.isclose(dperiod, dperiod[0]))

        if verbose:
            print('Frequencies set.')

    #--------------------------------------------------------------------------
    def set_tau(self, tau, verbose=True):
        """Set the taus (time points) for the analysis.

        Parameters
        ----------
        tau : array-like
            A range of time points at which the WWZ transform will be
            calculated.
        verbose : bool, optional
            Turns printing of information on or off. The default is True.

        Returns
        -------
        None.

        Notes
        -----
        It is recommended to use the get_tau() method to create the array of
        taus.
        """

        self.tau = np.asarray(tau)
        self.len0 = tau.size

        if verbose:
            print('Tau (time points) set.')

    #--------------------------------------------------------------------------
    def transform(self, c=1./(8.*np.pi**2), snr_weights=False, verbose=0):
        """Perform the WWZ transform as defined in [1].

        Parameters
        ----------
        c : float, optional
            The window decay parameter. See [1], Sec. 1 for a description.
            The default is 1./(8.*np.pi**2).
        snr_weights : bool, optional
            If True, the fluxes are additionally weighted by their
            corresponding signal-to-noise ratio, derived from the provided
            flux uncertainties. The default is False.
        verbose : int, optional
            Defines how much information is printed out. 0, no information is
            printed. 1, some information is printed. 2, all information is
            printed. The default is 0.

        Returns
        -------
        bool
            Returns True, if the analysis was performed succesfully. False,
            otherwise.

        References
        ----------
        [1] Foster, 1996
            https://ui.adsabs.harvard.edu/abs/1996AJ....112.1709F/abstract

        """

        if self.freq is None or self.tau is None:
            print('Frequencies and/or taus are not set. Analysis aborted.')
            return False

        if verbose:
            t_start = datetime.now()
            print('Starting the WWZ transform..')

        self.c = c
        self.snr_weights = snr_weights

        if verbose > 1:
            print('Setting frequency-tau-grid..')

        # create the wavelet grid:
        freq, tau = np.meshgrid(2*np.pi*self.freq, self.tau)
        # freq: the frequencies at which to evaluate the wavelet
        # tau: the delays over which to calculate the wavelet

        if verbose > 1:
            print('Creating weights..')

        # arrays which are used but not returned
        s = np.zeros((3, 3, self.len0, self.len1), dtype=np.float64)
        s_inv = np.zeros((3, 3, self.len0, self.len1), dtype=np.float64)
        w = np.zeros((self.n, self.len0, self.len1), dtype=np.float64)

        p1=0.0
        p2=0.0
        p3=0.0
        w_d=0.0
        w_d_s=0.0
        v_x=0.0
        v_x_s=0.0
        v_y=0.0
        v_y_s=0.0

        if verbose > 1:
            print('Calculating projections and scattering matrix..')

        for a in np.arange(0, self.n):
            # weights:
            w[a] = np.exp(-self.c * freq**2 * (self.t[a] - tau)**2)

            if snr_weights:
                w[a] *= np.sqrt(self.x[a]**2 / self.s_x[a]**2)

            w_d += w[a]
            w_d_s += w[a]**2

            # weighted variation of the data:
            v_x += w[a] * self.x[a]
            v_x_s += w[a] * self.x[a]**2

            # trial functions:
            t1 = 1.0
            t2 = np.cos(freq * (self.t[a] - tau))
            t3 = np.sin(freq * (self.t[a] - tau))

            # projections:
            p1 += w[a] * t1 * self.x[a]
            p2 += w[a] * t2 * self.x[a]
            p3 += w[a] * t3 * self.x[a]

            # s-matrix elements:
            s[0,0] += w[a] * t1 * t1
            s[1,0] += w[a] * t2 * t1
            s[0,1] += w[a] * t1 * t2
            s[2,0] += w[a] * t3 * t1
            s[0,2] += w[a] * t1 * t3
            s[1,1] += w[a] * t2 * t2
            s[2,1] += w[a] * t3 * t2
            s[1,2] += w[a] * t2 * t3
            s[2,2] += w[a] * t3 * t3

        # the projections of the trial functions onto the data are:
        p1 /= w_d
        p2 /= w_d
        p3 /= w_d
        # the scattering matrix elements are:
        s /= w_d

        if verbose > 1:
            print('Inverting scattering matrix..')

        # invert scattering matrix:
        for m in np.arange(0, self.len0):
            for n in np.arange(0, self.len1):
                s_inv[:,:,m,n] = np.linalg.inv(s[:,:,m,n])

        if verbose > 1:
            print('Calculating model coefficients..')

        # model coefficients:
        y1 = s_inv[0,0] * p1 + s_inv[0,1] * p2 + s_inv[0,2] * p3
        y2 = s_inv[1,0] * p1 + s_inv[1,1] * p2 + s_inv[1,2] * p3
        y3 = s_inv[2,0] * p1 + s_inv[2,1] * p2 + s_inv[2,2] * p3

        if verbose > 1:
            print('Calculating elements of transform..')

        # calculate elements of the transform:
        for a in np.arange(self.n):
            ya = y1 * 1.0 + y2 * np.cos(freq * (self.t[a] - tau)) \
                    + y3 * np.sin(freq * (self.t[a] - tau))
            v_y += w[a] * ya
            v_y_s += w[a] * ya**2

        if verbose > 1:
            print('Calculating WWZ and WWA..')

        self.v_x = (v_x_s / w_d) - (v_x / w_d)**2
        self.v_y = (v_y_s / w_d) - (v_y / w_d)**2
        self.n_eff = w_d**2 / w_d_s
        self.wwz = ((self.n_eff - 3.) * self.v_y) \
                / (2. * (self.v_x - self.v_y))
        self.wwa = np.sqrt(y2**2 + y3**2)

        if verbose:
            t_used = datetime.now() - t_start
            print('Finished in', t_used)

        return True

    #--------------------------------------------------------------------------
    def save(self, filename):
        """Save the class instance in a python pickle file.

        Parameters
        ----------
        filename : str
            Filename under which the class instance is saved.

        Returns
        -------
        None.
        """

        with open(filename, mode='wb') as f:
            pickle.dump(self, f)

    #--------------------------------------------------------------------------
    def load(self, filename):
        """Load a saved WWZ class instance from a python pickle file.

        Parameters
        ----------
        filename : str
            File name of the pickle file.

        Returns
        -------
        WWZ-type
            Returns a instance of WWZ.
        """

        with open(filename, mode='rb') as f:
            return pickle.load(f)

    #--------------------------------------------------------------------------
    def estimate_significance(self, simulations, append=False):
        """Monte Carlo based significance estimation.

        Parameters
        ----------
        simulations : list
            Each list entry is one simulation that will be analysed in the same
            way as the real data. The data structure for each simulation need
            to be the following. A list, tuple, or numpy.ndarray, where the
            first element/row contains the time steps, the second element/row
            contains the signal. Optionally a third element contains the
            corresponding uncertainties. The uncertainty can also be a single
            value.
        append : bool, optional
            If False, currently available significance estimations will be
            discarded and overwritten with the new analysis. If True, the
            analysis of the new simulations will be added to the existing ones.
            The default is False.

        Returns
        -------
        bool
            True, if the analysis finished. False, if the analysis was aborted.

        Notes
        -----
        When new simulations are appended, i.e. append=True, it is up to the
        user to ensure that the new simulations are independent of the previous
        ones. This method does not track simulations and cannot guarantee that
        repetitions are discarded.
        """

        if self.wwz is None:
            print('Analyse data first. Aborted!')
            return False

        # reset/initialize simulation results:
        if not append or self.wwz_sig is None:
            self.wwz_sig = np.zeros(shape=self.wwz.shape)
            self.wwa_sig = np.zeros(shape=self.wwa.shape)
            self.n_sig = 0

        # prepare WWZ instance for analysis of simulations:
        analyser = WWZ()

        n = len(simulations)
        t_start = datetime.now()
        print('Starting analysis of simulations at', t_start)

        # iterate through simulations:
        for i, simulation in enumerate(simulations):
            sys.stdout.write('\rProgress: {0:d} of {1:d}. {2:.1f}%'. format(
                    i+1, n, i*100./n))

            # extract simulation data:
            time = simulation[0]
            flux = simulation[1]
            try:
                flux_err = simulation[2]
            except:
                flux_err = 0

            # analyse simulation data:
            analyser.set_data(time, flux, flux_err, verbose=False)
            analyser.set_freq(self.freq, verbose=False)
            analyser.set_tau(self.tau, verbose=False)
            analyser.transform(
                    c=self.c, snr_weights=self.snr_weights, verbose=0)
            self.wwz_sig += analyser.wwz > self.wwz
            self.wwa_sig += analyser.wwa > self.wwa
            self.n_sig += 1

        self.wwz_pval = self.wwz_sig / self.n_sig
        self.wwa_pval = self.wwa_sig / self.n_sig

        t_used = datetime.now() - t_start
        print('\rFinished in', t_used)

        return True

    #--------------------------------------------------------------------------
    def _find_peaks(self, x, y, threshold, sampling=10):
        """Find peaks above a given threshold along a 1d-signal y(x).

        Parameters
        ----------
        x : numpy.ndarray
            Positional data of the signal, e.g. frequency or period.
        y : numpy.ndarray
            Signal, e.g. WWZ, WWA, or p-value along one time bin.
        threshold : float
            Defines the noise level. Only signals above this threshold are
            detected.
        sampling : int
            The interpolated signal is evalued at n times the sampling of the
            original data, where n is set by this number.

        Returns
        -------
        peak_x : list
            Peak positions.
        peak_y : list
            Peak signal strenghts.

        Notes
        -----
        This method uses spline interpolation to improve the estimation of the
        peak position and signal. Multiple peaks are identified recursively
        starting with the strongest peak.
        """

        if np.all(y < threshold) or x.size <= 3:
            return [], []

        # the spline fit requires an increasing order of x,
        # reverse order, if x is decreasing:
        if x.size > 1 and x[0] > x[1]:
            x = x[::-1]
            y = y[::-1]

        # spline interpolation:
        tck = splrep(x, y, s=0)
        n = x.size * int(sampling)
        x_interp = np.linspace(x[0], x[-1], n)
        y_interp = splev(x_interp, tck)
        y_der1 = splev(x_interp, tck, der=1)

        # identify peak:
        i = np.argmax(y_interp)

        # maximum at data edge, no peak found:
        if i + 1 == n or i == 0:
            return [], []

        # interpolate peak position:
        f = interp1d(y_der1[i-1:i+2], x_interp[i-1:i+2])
        peak_x = float(f(0))
        peak_y = float(splev(peak_x, tck))

        # find left peak edge:
        i = np.argmax(y)
        peak_max = peak_y
        n = x.size
        for n in range(i+1, x.size):
            if y[n] < peak_max:
                peak_max = y[n]
            else:
                break

        # find right peak edge:
        peak_max = peak_y
        m = 0
        for m in range(i-1, 0, -1):
            if y[m] < peak_max:
                peak_max = y[m]
            else:
                break

        peak_x = [peak_x]
        peak_y = [peak_y]

        # recursion on left part:
        if m > 1 and np.any(y[:m] > threshold):
            peak_x_l, peak_y_l = self._find_peaks(
                    x[:m], y[:m], threshold, sampling=sampling)
            peak_x = peak_x_l + peak_x
            peak_y = peak_y_l + peak_y

        # recursion on right part:
        if n < x.size-1 and np.any(y[n:] > threshold):
            peak_x_r, peak_y_r = self._find_peaks(
                    x[n:], y[n:], threshold, sampling=sampling)
            peak_x = peak_x + peak_x_r
            peak_y = peak_y + peak_y_r

        return peak_x, peak_y

    #--------------------------------------------------------------------------
    def find_peaks(
            self, signal, quantile, sampling=10, period_space=None,
            verbose=True):
        """Find peaks along frequency/period in a given signal for all time
        bins.

        Parameters
        ----------
        signal : string
            Select 'wwz' or 'wwa'.
        quantile : float
            Defines the noise level. Only signals above the threshold are
            detected that corresponds to the set quantile.
        sampling : int
            The interpolated signal is evalued at n times the sampling of the
            original data, where n is set by this number.
        verbose : bool, optional
            Turns printing of information on or off. The default is True.

        Raises
        ------
        ValueError
            Raised when 'signal' is not set to one of the allow values.

        Returns
        -------
        peak_tau : numpy.ndarray
            Time bins with detected signal above the threshold.
        peak_freq : numpy.ndarray
            Peak position, i.e. either frequency or period, depending on
            whether
            peak_ : numpy.ndarray

        Notes
        -----
        This method iterates through all time bins. The peak identification is
        implemented in the helper method self._find_peak().

        TBD
        ----
        Analysis of the p-values is currently turned off. The current peak
        location algorith does not work with flat p-value curves, due to an
        insufficient number of simulations.
        """

        # TODO: p-value analysis deactivated because it is running into bugs.

        peak_tau = []
        peak_pos = []
        peak_signal = []

        # select signal for analysis:
        if signal not in ['wwa', 'wwz']:
            raise ValueError(
                    "For 'signal' select 'wwz' or 'wwa'.")
        y = eval(f'self.{signal}')

        # invert p-values:
        if signal.find('pval') > -1:
            y = 1. - y

        # set threshold for signal detection:
        threshold = np.quantile(y, quantile)

        if verbose:
            print('Finding peaks in {0:s}.'.format(signal.upper()))
            print(f'Threshold set to: {threshold}')

        if (period_space is None and self.linear_period) or period_space:
            x = 1. / self.freq[::-1]
            y = y.transpose()[::-1].transpose()
            if verbose:
                print('Analysis in period space.')
        else:
            x = self.freq
            if verbose:
                print('Analysis in frequency space.')

        for i, t in enumerate(self.tau):
            peak_x, peak_y = self._find_peaks(
                    x, y[i], threshold, sampling=sampling)
            for px, py in zip(peak_x, peak_y):
                peak_tau.append(t)
                peak_pos.append(px)
                peak_signal.append(py)

        peak_tau = np.array(peak_tau)
        peak_pos = np.array(peak_pos)
        peak_signal = np.array(peak_signal)

        if verbose:
            print('Done.')

        return peak_tau, peak_pos, peak_signal

#==============================================================================


from unittest import mock

from idb.daemon.server import CompositeServer
from idb.utils.testing import TestCase, ignoreTaskLeaks

@ignoreTaskLeaks
class ServerTest(TestCase):
    def test_close(self) -> None:
        first = mock.MagicMock()
        second = mock.MagicMock()
        server = CompositeServer(servers=[first, second], logger=mock.MagicMock())
        server.close()
        first.close.assert_called_once()
        second.close.assert_called_once()


import random
import logging
import os
import importlib.util
from pathlib import Path
from contextlib import ExitStack

import numpy as np
import boto3
from botocore.exceptions import ClientError

from ..utils.factory import get_monitor, get_model
from ..utils.aws import s3_fetch_data
from ..utils.misc import set_framework_seed

logging.basicConfig(level=logging.INFO)

def main() -> None:
    s3_path = Path(os.environ['S3_PATH'])
    bucket_name = os.environ['BUCKETNAME']
    blueprint_path = Path('home', 'dl-benchmarking', 'bp.py')

    s3 = boto3.resource('s3')
    bucket = s3.Bucket(bucket_name)
    bucket.download_file(str(s3_path / 'bp.py'), str(blueprint_path))

    spec = importlib.util.spec_from_file_location('bp', blueprint_path)
    bp = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(bp)

    logging.info(f'{bp.framework}_{bp.job_name}_{bp.instance_type}')

    random.seed(bp.python_seed)
    np.random.seed(bp.numpy_seed)
    set_framework_seed(bp.framework, bp.framework_seed)

    train_data = s3_fetch_data(bp.bucket_name,
                               Path(bp.train_data_path)).reshape(*bp.reshape_input_data)
    train_label = s3_fetch_data(bp.bucket_name,
                                Path(bp.train_label_path))

    try:
        test_data = s3_fetch_data(bp.bucket_name,
                                  Path(bp.test_data_path)).reshape(*bp.reshape_input_data)
        test_label = s3_fetch_data(bp.bucket_name,
                                   Path(bp.test_label_path))
    except (TypeError, ClientError):
        test_data = None
        test_label = None

    try:
        params = s3_fetch_data(bucket_name, Path(bp.model_params_path))
    except (TypeError, ClientError):
        params = None

    if bp.framework in ['pytorch', 'mxnet']: # Use NCHW format for image data in pytorch and mxnet
        if len(train_data.shape) > 3:
            train_data = train_data.swapaxes(-3, -1)
            if isinstance(test_data, np.ndarray):
                test_data = test_data.swapaxes(-3, -1)

    model = get_model(bp.architecture, bp.framework, params)

    if bp.mode == 'training':
        with ExitStack() as stack:
            for monitor in bp.monitors:
                stack.enter_context(get_monitor(monitor, s3_path, bp.bucket_name, aws=True))
            model.training(train_data,
                           train_label,
                           test_data,
                           test_label,
                           bp.loss_function,
                           bp.optimizer,
                           bp.optimizer_mode,
                           bp.epochs,
                           bp.learning_rate,
                           s3_path,
                           bp.batch_size,
                           bp.train_seq_length,
                           bp.test_seq_length,
                           bp.state_sizes,
                           bp.truncated_backprop_length,
                           bp.bucket_name,
                           aws=True)

    elif bp.mode == 'inference':
        with ExitStack() as stack:
            for monitor in bp.monitors:
                stack.enter_context(get_monitor(monitor, s3_path, bp.bucket_name, aws=True))
            model.inference(train_data,
                            bp.batch_size,
                            s3_path,
                            bp.mode,
                            bp.optimizer_mode,
                            bp.state_sizes,
                            bp.bucket_name,
                            aws=True)

if __name__ == "__main__":
    main()


def rate_trader_balances(row):
    # How much BASE plus RAI in BASE is held
    traders = row['rate_traders']
    rp = row['target_price']

    total_base = 0
    base = 0
    total_rai_base = 0
    for trader in traders:
        total_base += trader['base_balance']
        total_base += trader['rai_balance'] * rp
        base += trader['base_balance']
        total_rai_base += trader['rai_balance'] * rp

    return total_base, base, total_rai_base

def price_trader_balances(row):
    # How much BASE plus RAI in BASE is held
    traders = row['price_traders']
    rp = row['target_price']

    total_base = 0
    base = 0
    total_rai_base = 0
    for trader in traders:
        total_base += trader['base_balance']
        total_base += trader['rai_balance'] * rp
        base += trader['base_balance']
        total_rai_base += trader['rai_balance'] * rp

    return total_base, base, total_rai_base

def eth_leverager_balances(row):
    eth_price = row['eth_price']
    rp = row['target_price']
    safes = row['cdps'] # SAFEs are initialized in timestep 1

    df = safes.query(f"open == 1 and owner == 'leverager'")
    total_collateral = df['locked'].sum() - df['freed'].sum() - df['v_bitten'].sum()
    total_debt = df['drawn'].sum() - df['wiped'].sum() - df['u_bitten'].sum()
    debt_base = total_debt * rp
    collateral_base = total_collateral * eth_price
    total_base = collateral_base - debt_base

    return total_base, collateral_base, debt_base, total_collateral, total_debt

"""
def eth_leverager_total_base(row):
    eth_price = row['eth_price']
    rp = row['target_price']
    safes = row['cdps'] # SAFEs are initialized in timestep 1

    collateral_base = 0
    debt_base = 0
    for index, cdp in safes.query(f"open == 1 and owner == 'leverager'").iterrows():
        locked = cdp["locked"]
        freed = cdp["freed"]
        drawn = cdp["drawn"]
        wiped = cdp["wiped"]
        dripped = cdp["dripped"]
        v_bitten = cdp["v_bitten"]
        u_bitten = cdp["u_bitten"]
        w_bitten = cdp["w_bitten"]

        collateral = locked - freed - v_bitten
        debt = drawn - wiped - u_bitten

        collateral_base += collateral * eth_price
        debt_base += debt * rp

    return collateral_base - debt_base

def eth_leverager_collateral_base(row):
    eth_price = row['eth_price']
    rp = row['target_price']
    safes = row['cdps'] # SAFEs are initialized in timestep 1

    collateral_base = 0
    for index, cdp in safes.query(f"open == 1 and owner == 'leverager'").iterrows():
        locked = cdp["locked"]
        freed = cdp["freed"]
        drawn = cdp["drawn"]
        wiped = cdp["wiped"]
        dripped = cdp["dripped"]
        v_bitten = cdp["v_bitten"]
        u_bitten = cdp["u_bitten"]
        w_bitten = cdp["w_bitten"]

        collateral = locked - freed - v_bitten
        collateral_base += collateral * eth_price

    return collateral_base

def eth_leverager_collateral(row):
    eth_price = row['eth_price']
    rp = row['target_price']
    safes = row['cdps'] # SAFEs are initialized in timestep 1

    collateral = 0
    for index, cdp in safes.query(f"open == 1 and owner == 'leverager'").iterrows():
        locked = cdp["locked"]
        freed = cdp["freed"]
        drawn = cdp["drawn"]
        wiped = cdp["wiped"]
        dripped = cdp["dripped"]
        v_bitten = cdp["v_bitten"]
        u_bitten = cdp["u_bitten"]
        w_bitten = cdp["w_bitten"]

        collateral = locked - freed - v_bitten
        collateral += collateral

    return collateral

def eth_leverager_debt_base(row):
    rp = row['target_price']
    safes = row['cdps'] # SAFEs are initialized in timestep 1

    debt_base = 0
    for index, cdp in safes.query(f"open == 1 and owner == 'leverager'").iterrows():
        locked = cdp["locked"]
        freed = cdp["freed"]
        drawn = cdp["drawn"]
        wiped = cdp["wiped"]
        dripped = cdp["dripped"]
        v_bitten = cdp["v_bitten"]
        u_bitten = cdp["u_bitten"]
        w_bitten = cdp["w_bitten"]

        debt = drawn - wiped - u_bitten
        debt_base += debt * rp

    return debt_base

def eth_leverager_debt(row):
    safes = row['cdps'] # SAFEs are initialized in timestep 1

    debt = 0
    for index, cdp in safes.query(f"open == 1 and owner == 'leverager'").iterrows():
        locked = cdp["locked"]
        freed = cdp["freed"]
        drawn = cdp["drawn"]
        wiped = cdp["wiped"]
        dripped = cdp["dripped"]
        v_bitten = cdp["v_bitten"]
        u_bitten = cdp["u_bitten"]
        w_bitten = cdp["w_bitten"]

        debt += drawn - wiped - u_bitten

    return debt

"""
def price_trader_profit(df, run):
    df_run = df.query(f'run == {run}')
    start_traders = df_run['price_traders'].iloc[1]
    start_rp = df_run['target_price'].iloc[1]

    end_traders = df_run['price_traders'].iloc[-1]
    end_rp = df_run['target_price'].iloc[-1]

    start_total_base = 0
    for trader in start_traders:
        start_total_base += trader['base_balance']
        start_total_base += trader['rai_balance'] * start_rp

    end_total_base = 0
    for trader in end_traders:
        end_total_base += trader['base_balance']
        end_total_base += trader['rai_balance'] * end_rp

    return end_total_base - start_total_base

def rate_trader_profit(df, run):
    df_run = df.query(f'run == {run}')
    start_traders = df_run['rate_traders'].iloc[1]
    start_rp = df_run['target_price'].iloc[1]

    final_traders = df_run['rate_traders'].iloc[-1]
    final_rp = df_run['target_price'].iloc[-1]

    start_total_base = 0
    start_total_rai = 0
    start_total_rai_base = 0
    start_base_only = 0
    for trader in start_traders:
        start_base_only += trader['base_balance']
        start_total_base += trader['base_balance']
        start_total_base += trader['rai_balance'] * start_rp
        start_total_rai += trader['rai_balance']
        start_total_rai_base += trader['rai_balance'] * start_rp

    final_total_base = 0
    final_total_rai = 0
    final_total_rai_base = 0
    final_base_only = 0
    n_trades = 0
    for trader in final_traders:
        final_base_only += trader['base_balance']
        final_total_base += trader['base_balance']
        final_total_base += trader['rai_balance'] * final_rp
        final_total_rai += trader['rai_balance']
        final_total_rai_base += trader['rai_balance'] * final_rp
        n_trades += trader['n_buys'] + trader['n_sells']

    return {'start_base_only': start_base_only,
            'start_total_rai': start_total_rai,
            'start_total_rai_base': start_total_rai_base,
            'start_total_base': start_total_base,
            'final_base_only': final_base_only,
            'final_total_rai': final_total_rai,
            'final_total_rai_base': final_total_rai_base,
            'final_total_base': final_total_base,
            'n_trades': n_trades,
            'profit': final_total_base - start_total_base,
            'profit_per_trade': (final_total_base - start_total_base)/n_trades

           }

def cdp_profit(df, owner, run):
    df_run = df.query(f'run == {run}')
    start_eth_price = df_run['eth_price'].iloc[0]
    start_rp = df_run['target_price'].iloc[0]
    start_safes = df_run['cdps'].iloc[1] # SAFEs are initialized in timestep 1

    final_eth_price = df_run['eth_price'].iloc[-1]
    final_rp = df_run['target_price'].iloc[-1]
    final_safes = df_run['cdps'].iloc[-1]

    start_base = 0
    start_collateral = 0
    start_collateral_base = 0
    start_debt = 0
    start_debt_base = 0
    for index, cdp in start_safes.query(f"open == 1 and owner == '{owner}'").iterrows():
        locked = cdp["locked"]
        freed = cdp["freed"]
        drawn = cdp["drawn"]
        wiped = cdp["wiped"]
        dripped = cdp["dripped"]
        v_bitten = cdp["v_bitten"]
        u_bitten = cdp["u_bitten"]
        w_bitten = cdp["w_bitten"]

        collateral = locked - freed - v_bitten
        debt = drawn - wiped - u_bitten

        start_collateral += collateral
        start_debt += debt

        collateral_base = collateral * start_eth_price
        debt_base = debt * start_rp

        #print(f"start {collateral_base=}, {collateral=}, {start_eth_price=}")

        start_collateral_base += collateral_base
        start_debt_base += debt_base
        start_base += collateral_base - debt_base

    final_base = 0
    final_collateral = 0
    final_collateral_base = 0
    final_debt = 0
    final_debt_base = 0
    for index, cdp in final_safes.query(f"open == 1 and owner == '{owner}'").iterrows():
        locked = cdp["locked"]
        freed = cdp["freed"]
        drawn = cdp["drawn"]
        wiped = cdp["wiped"]
        dripped = cdp["dripped"]
        v_bitten = cdp["v_bitten"]
        u_bitten = cdp["u_bitten"]
        w_bitten = cdp["w_bitten"]

        collateral = locked - freed - v_bitten
        debt = drawn - wiped - u_bitten

        final_collateral += collateral
        final_debt += debt
        collateral_base = collateral * final_eth_price
        debt_base = debt * final_rp

        #print(f"final {collateral_base=}, {collateral=}, {final_eth_price=}")
        final_collateral_base += collateral_base
        final_debt_base += debt_base

        final_base += collateral_base - debt_base

    return {"start_collateral": start_collateral,
            "start_collateral_base": start_collateral_base,
            "start_debt": start_debt,
            "start_debt_base": start_debt_base,
            "final_collateral": final_collateral,
            "final_collateral_base": final_collateral_base,
            "final_debt": final_debt,
            "final_debt_base": final_debt_base,
            "final_new_debt": final_debt - start_debt,
            "final_new_debt_base": final_debt_base - start_debt_base,
            "start_base": start_base,
            "final_base": final_base,
            "profit": final_collateral_base - start_collateral_base - final_debt_base + start_debt_base
           }


"""A module containing the agents which can learn strategies for workers."""

from os import path
import numpy as np
from numpy import random
import matplotlib.pyplot as plt
from matplotlib import colors, ticker

import environment.bitmanipulation as bitm

###############################################################################
# Agent
###############################################################################
#
class QLearningAgent:
    """Agent which employs Q learning to make decisions."""
    def __init__(self, config, config_dir, get_action_num_func):
        # settings
        self.name = config["name"]
        self.config_dir = config_dir
        self.deadline = config["deadline"]

        # learning variables
        self.learning_rate = config["learning rate"]
        self.discount = config["discount factor"]

        # set up actions
        self._possible_actions_str = config["possible actions"]
        self.possible_actions = []
        for action_name in self._possible_actions_str:
            # convert the actions names to numbers, with the given function
            self.possible_actions.append(get_action_num_func(action_name))

        # this maps the external action number to internal action indices
        self.action_num2idx = {}
        for action_idx, action_num in enumerate(self.possible_actions):
            self.action_num2idx[action_num] = action_idx

        self.state_space_type = config["state space"]["type"]
        if self.state_space_type == "time fitness":
            self.quantisation_levels = \
                    config["state space"]["quantisation levels"]

            self._find_state = self._find_state_time_fitness

            q_table_shape = [
                self.deadline +1, # time
                self.quantisation_levels, # fitness
                len(self.possible_actions), # possible actions
            ]
        elif self.state_space_type == "time memory":
            assert len(self.possible_actions) == 2
            self.history = \
                    config["state space"]["history"]

            self._find_state = self._find_state_time_memory

            q_table_shape = [
                self.deadline +1, # time
                1 << self.history, # memory
                len(self.possible_actions), # possible actions
            ]
        else:
            raise ValueError("State space type not valid")

        self._q_table = np.zeros(q_table_shape)

        # the count of the number of times each state has been updated
        self._update_count = np.zeros(q_table_shape[:-1])

        # exploration type
        if config["exploration"]["type"] == "epsilon greedy":
            self.epsilon = config["exploration"]["epsilon start"]
            self.epsilon_decay = config["exploration"]["epsilon decay"]

            self.best_action = self._choose_greedy_action
            self.explore_action = self._choose_epsilon_greedy_action
            self.episode_end = self._decay_epsilon

        elif config["exploration"]["type"] == "boltzmann":
            self.best_action = self._choose_greedy_action
            self.explore_action = self._choose_boltzmann_action
            self.episode_end = lambda *args : None

        else:
            raise ValueError("Exploration type not valid")

    def _find_state_time_fitness(self, node, time, environment):
        """Find a worker's state at the given time.
        The state space used is the time and the quantised normalised fitness.
        """
        current_fitness = environment.get_node_fitness_norm(node, time)
        return (time, int(current_fitness * (self.quantisation_levels -1)))

    def _find_state_time_memory(self, node, time, environment):
        """Find a worker's state at the given time.
        The state space used is the time and the binary memory.

        Binary memory is an integer with bit-length of the given history.

        The value of each bit is the action
        which was taken in a given time-step of the last few time-steps.

        For example, the number 11 (1011 in binary),
        when there is a history of 4, would mean
        that action 1 was taken taken at time -1, time -2 and time -4,
        and so action 0 was only taken at time -3.

        | time -4 | time -3 | time -2 | time -1 |
        |    1    |    0    |    1    |    1    |
        """
        end = time if time < self.history else self.history

        memory = 0
        for idx in range(end):
            action_idx = self.action_num2idx[
                    environment.get_node_action(node, time - idx - 1)
            ]
            memory = bitm.set_bit(memory, idx, action_idx)

        return (time, memory)

    def _update_q_table(self, prior_state, action_num, post_state, reward):
        """Updates the Q table using the given transition."""
        # the action index of the best action in the post transition state.
        post_best_action_idx = np.argmax(self._q_table[post_state])
        td_target = reward \
            + self.discount * self._q_table[post_state][post_best_action_idx]

        action_idx = self.action_num2idx[action_num]
        td_delta = td_target - self._q_table[prior_state][action_idx]

        self._q_table[prior_state][action_idx] += self.learning_rate * td_delta
        self._update_count[prior_state] += 1

    def _choose_greedy_action(self, node, time, environment):
        """Returns the best action according the Q table."""
        current_state = self._find_state(node, time, environment)
        return self.possible_actions[np.argmax(self._q_table[current_state])]

    def _choose_epsilon_greedy_action(self, node, time, environment):
        """Returns either the best action according the Q table
        or a random action depending on the current epsilon value.
        """
        current_state = self._find_state(node, time, environment)
        # choose action
        if random.rand() <= self.epsilon:
            action = random.choice(self.possible_actions)
        else:
            action = self.possible_actions[
                np.argmax(self._q_table[current_state])
            ]
        return action

    def _choose_boltzmann_action(self, node, time, environment):
        """Returns either the action using the softmax
        as a probability distribution"""
        current_state = self._find_state(node, time, environment)
        action_values = self._q_table[current_state]

        action_values_exp = \
                np.array([np.exp(action) for action in action_values])

        probs = action_values_exp / np.sum(action_values_exp)

        action_num = np.random.choice(
                self.possible_actions, 1, replace=False, p=probs,
        )
        return action_num

    def _decay_epsilon(self):
        """Decays the value of epsilon by one unit."""
        self.epsilon -= self.epsilon_decay * self.epsilon

    def learn(self, node, prior_time, post_time, environment):
        """Learns from a transition."""
        prior_state = self._find_state(node, prior_time, environment)
        post_state = self._find_state(node, post_time, environment)

        # each node only receives a reward at the deadline
        reward = environment.get_node_fitness(node, post_time) \
                if post_time == self.deadline else 0

        chosen_action = environment.get_node_action(node, prior_time)
        self._update_q_table(prior_state, chosen_action, post_state, reward)

    def _find_file_name(self, suffix):
        """Finds the file name for a given suffix."""
        if suffix:
            return path.join(self.config_dir, f'{self.name}_{suffix}.npz')
        return path.join(self.config_dir, f'{self.name}.npz')

    def save(self, suffix=None):
        """Saves the agent's learnings.
        Saves the q_table and update_count in a file with the the given suffix.
        """
        file_name = self._find_file_name(suffix)
        with open(file_name, 'wb') as file_handle:
            np.savez(
                file_handle,
                q_table=self._q_table,
                update_count=self._update_count,
            )

    def load(self, suffix=None):
        """Loads learnings.
        Loads the q_table and update_count from a file with the the given
        suffix.
        """
        file_name = self._find_file_name(suffix)
        with open(file_name, 'rb') as file_handle:
            file_content = np.load(file_handle)
            self._q_table = file_content["q_table"]
            self._update_count = file_content["update_count"]

    def plot(self, figure):
        """Plots the q table and update count.
        Plots the q table and update count of an agent in the provided figure.
        """
        if self.state_space_type == "time memory":
            len_bin_labels = self.history
            y_stride = 1
            print(
                "Binary memory values:\n\t0 -", self._possible_actions_str[0],
                "\n\t1 -", self._possible_actions_str[1],
            )
            ylabel = "Binary Memory"
        else:
            len_bin_labels = None
            y_stride = 2
            ylabel = "Fitness"

        num_possible_actions = len(self.possible_actions)
        if num_possible_actions == 2:
            axs = figure.subplots(3, 1)

            plot_q_table_image(
                self._possible_actions_str,
                self._q_table,
                axis=axs[0],
                ytick_stride=y_stride,
                len_binary_yticklabels=len_bin_labels,
            )
            plot_q_table_action_image(
                self._possible_actions_str,
                0,
                self._q_table,
                axis=axs[1],
                ytick_stride=y_stride,
                len_binary_yticklabels=len_bin_labels,
            )
            plot_update_count_image(
                self._update_count,
                axis=axs[2],
                ytick_stride=y_stride,
                len_binary_yticklabels=len_bin_labels,
            )

            for axis in axs:
                axis.set_ylabel(ylabel)
                axis.set_xlabel("Time Step")

        elif num_possible_actions == 3:

            axs = figure.subplots(3, 2)
            plot_q_table_image(
                self._possible_actions_str, self._q_table, axis=axs[0, 0]
            )
            plot_q_table_action_image(
                self._possible_actions_str, 0, self._q_table, axis=axs[1, 0]
            )
            plot_q_table_action_image(
                self._possible_actions_str, 1, self._q_table, axis=axs[2, 0]
            )
            plot_update_count_image(self._update_count, axis=axs[0, 1])
            plot_q_table_action_image(
                self._possible_actions_str, 2, self._q_table, axis=axs[1, 1]
            )
            axs[2, 1].remove()

            for i in range(3):
                for j in range(2):
                    axs[i, j].set_ylabel("Fitness")
                    axs[i, j].set_xlabel("Time Step")

        else:
            raise ValueError("Can only plot tables for agents with"
                             " two or three possible actions."
                             " This agent has {num_possible_actions}.")

        plt.tight_layout()

###############################################################################
# Functions
###############################################################################
#
def plot_update_count_image(
        update_count,
        axis=None,
        ytick_stride=2,
        xtick_stride=2,
        len_binary_yticklabels=None,
):
    """Plots the given update_count as an image."""
    if not axis:
        axis = plt.gca()

    update_count = np.swapaxes(update_count, 0, 1)
    update_count = update_count[:, :-1]

    image = axis.imshow(update_count, norm=colors.LogNorm())
    cbar = axis.figure.colorbar(image, ax=axis)
    cbar.ax.set_ylabel("Number of Updates")

    # disable splines
    for _, spine in axis.spines.items():
        spine.set_visible(False)

    axis.set_xticks(np.arange(0.5, update_count.shape[1], 1), minor=True)
    axis.set_yticks(np.arange(0.5, update_count.shape[0], 1), minor=True)

    yticks = np.arange(0, update_count.shape[0], ytick_stride)
    axis.set_yticks(yticks)
    axis.set_xticks(np.arange(0, update_count.shape[1], xtick_stride))

    if len_binary_yticklabels:
        axis.set_yticklabels(
                [format(i, f'0{len_binary_yticklabels}b') for i in yticks]
        )

    axis.set_xlim(-0.5, update_count.shape[1])
    axis.set_ylim(-0.5, update_count.shape[0])

    axis.grid(which="minor", color="w", linestyle="-", linewidth=2)
    axis.tick_params(which="minor", bottom=False, left=False)

    axis.set_title("Experience")

def plot_q_table_image(
        possible_actions,
        q_table,
        axis=None,
        ytick_stride=2,
        xtick_stride=2,
        len_binary_yticklabels=None,
):
    """Plots the preferred action in each state for a given q table."""

    if not axis:
        axis = plt.gca()

    # the actions with the highest perceived value
    best_actions = np.argmax(q_table, axis=2)

    best_actions = np.swapaxes(best_actions, 0, 1)
    best_actions = best_actions[:, :-1]

    colour_tuple = ("g", "b", "m")

    fmt = ticker.FixedFormatter(possible_actions)
    cbar_ticks_len = len(possible_actions)
    cbar_ticks = np.arange(cbar_ticks_len, dtype="f")

    if cbar_ticks_len == 3:
        cbar_ticks[0] = 1./3
        cbar_ticks[cbar_ticks_len -1] = cbar_ticks_len - 1 - 1./3
    elif cbar_ticks_len == 2:
        cbar_ticks[0] = 1./4
        cbar_ticks[cbar_ticks_len -1] = cbar_ticks_len - 1 - 1./4

    cmap = colors.ListedColormap(colour_tuple[:cbar_ticks_len])

    image = axis.imshow(best_actions, cmap=cmap)
    axis.figure.colorbar(image, ax=axis, format=fmt, ticks=cbar_ticks)

    # disable splines
    for _, spine in axis.spines.items():
        spine.set_visible(False)

    axis.set_xticks(np.arange(0.5, best_actions.shape[1], 1), minor=True)
    axis.set_yticks(np.arange(0.5, best_actions.shape[0], 1), minor=True)

    yticks = np.arange(0, best_actions.shape[0], ytick_stride)
    axis.set_yticks(yticks)
    axis.set_xticks(np.arange(0, best_actions.shape[1], xtick_stride))

    if len_binary_yticklabels:
        axis.set_yticklabels(
                [format(i, f'0{len_binary_yticklabels}b') for i in yticks]
        )

    axis.set_xlim(-0.5, best_actions.shape[1])
    axis.set_ylim(-0.5, best_actions.shape[0])

    axis.grid(which="minor", color="w", linestyle="-", linewidth=2)
    axis.tick_params(which="minor", bottom=False, left=False)

    axis.set_title("Learnt Policy")

def plot_q_table_action_image(
        possible_actions,
        action_idx,
        q_table,
        axis=None,
        ytick_stride=2,
        xtick_stride=2,
        len_binary_yticklabels=None,
        colour_map="plasma",
):
    """Plots the relative value of an action at each time step."""

    if not axis:
        axis = plt.gca()

    diff_actions = q_table[:, :, action_idx]/np.sum(q_table, axis=2)

    diff_actions = np.swapaxes(diff_actions, 0, 1)
    diff_actions = diff_actions[:, :-1]

    image = axis.imshow(diff_actions, cmap=colour_map, vmin=0, vmax=1)
    cbar = axis.figure.colorbar(image, ax=axis)
    cbar.ax.set_ylabel("Relative Value of Action")

    # disable splines
    for _, spine in axis.spines.items():
        spine.set_visible(False)

    axis.set_xticks(np.arange(0.5, diff_actions.shape[1], 1), minor=True)
    axis.set_yticks(np.arange(0.5, diff_actions.shape[0], 1), minor=True)

    yticks = np.arange(0, diff_actions.shape[0], ytick_stride)
    axis.set_yticks(yticks)
    axis.set_xticks(np.arange(0, diff_actions.shape[1], xtick_stride))

    if len_binary_yticklabels:
        axis.set_yticklabels(
                [format(i, f'0{len_binary_yticklabels}b') for i in yticks]
        )

    axis.set_xlim(-0.5, diff_actions.shape[1])
    axis.set_ylim(-0.5, diff_actions.shape[0])

    axis.grid(which="minor", color="w", linestyle="-", linewidth=2)
    axis.tick_params(which="minor", bottom=False, left=False)

    axis.set_title(f"Relative Value of '{possible_actions[action_idx]}'")


#!/usr/bin/env python
import numpy as np
import random
import cv2

def do_full_homo(p, q):
    mat = np.zeros((len(p) * 3, 9))
    for i in range(len(p)):
        mat[i * 2][0] = p[i][0]
        mat[i * 2][1] = p[i][1]
        mat[i * 2][2] = 1

        mat[i * 2][3] = 0
        mat[i * 2][4] = 0
        mat[i * 2][5] = 0

        mat[i * 2][6] = -p[i][0] * q[i][0]
        mat[i * 2][7] = -p[i][1] * q[i][0]
        mat[i * 2][8] = -q[i][0]

        mat[(i * 2) + 1][0] = 0
        mat[(i * 2) + 1][1] = 0
        mat[(i * 2) + 1][2] = 0

        mat[(i * 2) + 1][3] = p[i][0]
        mat[(i * 2) + 1][4] = p[i][1]
        mat[(i * 2) + 1][5] = 1

        mat[(i * 2) + 1][6] = -p[i][0] * q[i][1]
        mat[(i * 2) + 1][7] = -p[i][1] * q[i][1]
        mat[(i * 2) + 1][8] = -q[i][1]
    _, _, v = np.linalg.svd(mat)
    h = v[-1].reshape(3, 3)
    h /= v[-1][-1]
    return h

def ransac(points_left, points_right, desc_left, desc_right):
    n = 10000
    best_points = 0
    best_homo = None

    max_len = min(len(points_left), len(points_right))

    bf = cv2.BFMatcher()
    matches = bf.match(desc_left, desc_right)
    mp1 = [points_left[m.queryIdx].pt for m in matches]
    mp2 = [points_right[m.trainIdx].pt for m in matches]
    #print(mp1, mp2)

    for _ in range(n):
        l = [random.randint(0, max_len - 1) for _ in range(4)]
        p = np.array([mp1[i] for i in l])
        q = np.array([mp2[i] for i in l])

        #potential_homo = do_full_homo(p, q)
        potential_homo, _ = cv2.findHomography(p, q)
        if np.array_equal(potential_homo, None):
            continue
        points_transformed = [potential_homo @ (point[0], point[1], 1) for point in mp1]
        good_points = 0
        for i in range(len(mp1)):
            p_i = points_transformed[i][:2]
            q_i = mp2[i]
            #print(mp1[i], p_i, q_i)
            diff = np.linalg.norm(p_i - q_i)
            #if diff < 100:
                #print(diff)
            if diff < 10:
                good_points += 1
            if best_points < good_points:
                best_points = good_points
                best_homo = potential_homo
    #print(best_points, "out of", len(points_left))
    return best_homo

    """
    for _ in range(n):
        l = [random.randint(0, max_len - 1) for _ in range(4)]
        p = [points_left[i].pt for i in l]
        q = [points_right[i].pt for i in l]

        potential_homo = do_full_homo(p, q)
        points_transformed = [potential_homo @ (point.pt[0], point.pt[1], 1) for point in points_left]
        good_points = 0
        for i in range(max_len):
            p_i = points_transformed[i][:2]
            q_i = points_right[i].pt
            diff = np.linalg.norm(p_i - q_i)
            #total_diff += diff
            if diff < 50:
                good_points += 1
        if best_points < good_points:
            best_points = good_points
            best_homo = potential_homo
    print(best_points, "out of", len(points_left))
    return best_homo
    """

path_left = './to_detect_2_1.jpg'
path_right = './to_detect_2_2.jpg'

img_left = cv2.imread(path_left)
img_right = cv2.imread(path_right)

sift = cv2.SIFT_create()
points_left, desc_left = sift.detectAndCompute(img_left, None)
points_right, desc_right = sift.detectAndCompute(img_right, None)

homo = ransac(points_left, points_right, desc_left, desc_right)

#max_len = min(len(points_left), len(points_right))
#pts_src = np.array([np.float32(k.pt) for k in points_left], np.float32).reshape(-1, 2)[:max_len]
#pts_dst = np.array([np.float32(k.pt) for k in points_right], np.float32).reshape(-1, 2)[:max_len]
#other_homo, _ = cv2.findHomography(pts_src, pts_dst)

warped = cv2.warpPerspective(img_right, np.linalg.inv(homo), (img_right.shape[1] * 2, img_right.shape[0]))
warped[:img_left.shape[0], :img_left.shape[1]] = img_left
cv2.imwrite("stitched.jpg", warped)


from db import new_connection
from mysql.connector.errors import ProgrammingError

sql='INSERT INTO contatos (nome,telefone) values(%s,%s)'
args=(('lucas','34233-4123'),
    ('alex','34233-4123'),
    ('bruno','34233-4123'),
    ('maria','34233-4123'),
    ('valentina','34233-4123'),
    ('carol','34233-4123'),
    ('beca','34233-4123'),)

with new_connection() as connection:
    try:
        cursor=connection.cursor()
        cursor.executemany(sql,args)
        connection.commit()
    except ProgrammingError as e:
        print(f'ERROR: {e}')


"""
2019 Day 14
https://adventofcode.com/2019/day/14
"""

from functools import partial
from math import ceil, floor
from typing import Callable, Dict, List, Set, Tuple
import aocd  # type: ignore

MaterialQty = Tuple[str, int]
Recipes = Dict[str, Tuple[int, Set[MaterialQty]]]

def read_ingredient(ingredient: str) -> MaterialQty:
    parts = ingredient.split(" ")
    return parts[1], int(parts[0])

def read_recipe_side(side: str) -> List[MaterialQty]:
    return [read_ingredient(ing) for ing in side.split(", ")]

def read_recipe_line(line: str) -> Tuple[List[MaterialQty], List[MaterialQty]]:
    sides = [read_recipe_side(side) for side in line.split(" => ")]
    return sides[0], sides[1]

def read_recipes(text: str) -> Recipes:
    return {
        right[0][0]: (right[0][1], set(left))
        for left, right in [read_recipe_line(line) for line in text.split("\n")]
    }

def resolve_need(
    element: str, inventory: Dict[str, int], recipes: Recipes
) -> Dict[str, int]:
    reacted = dict(inventory)
    need = -inventory[element]

    if need > 0 and element in recipes:
        qty, ingredients = recipes[element]
        repetitions = ceil(need / qty)
        reacted[element] = reacted.get(element, 0) + (repetitions * qty)
        for (ing_element, ing_qty) in ingredients:
            reacted[ing_element] = reacted.get(ing_element, 0) - (repetitions * ing_qty)

    return reacted

def ore_required(recipes: Recipes, fuel: int = 1) -> int:
    inventory = {"FUEL": -fuel}
    needs = ["FUEL"]
    while len(needs) > 0:
        inventory = resolve_need(needs[0], inventory, recipes)
        needs = [
            element
            for element, qty in inventory.items()
            if element != "ORE" and qty < 0
        ]
    return -inventory["ORE"]

def binary_search(
    func: Callable[[int], int], target: int, lower: int, upper: int
) -> int:
    while lower <= upper:
        mid = floor((lower + upper) / 2)
        attempt = func(mid)
        if attempt < target:
            lower = mid + 1
        elif attempt > target:
            upper = mid - 1
        else:
            return mid
    return upper

def exponential_search(func: Callable[[int], int], target: int) -> int:
    bound = 1
    while func(bound) < target:
        bound *= 2
    return binary_search(func, target, bound // 2, bound + 1)

def fuel_from_trillion_ore(recipes: Recipes) -> int:
    return exponential_search(partial(ore_required, recipes), 1_000_000_000_000)

def test_part1() -> None:
    """
    Examples for Part 1.
    """
    text = "\n".join(
        (
            "10 ORE => 10 A",
            "1 ORE => 1 B",
            "7 A, 1 B => 1 C",
            "7 A, 1 C => 1 D",
            "7 A, 1 D => 1 E",
            "7 A, 1 E => 1 FUEL",
        )
    )
    eg1 = {
        "A": (10, {("ORE", 10)}),
        "B": (1, {("ORE", 1)}),
        "C": (1, {("A", 7), ("B", 1)}),
        "D": (1, {("A", 7), ("C", 1)}),
        "E": (1, {("A", 7), ("D", 1)}),
        "FUEL": (1, {("A", 7), ("E", 1)}),
    }
    assert read_recipes(text) == eg1
    assert ore_required(eg1) == 31
    eg2 = {
        "A": (2, {("ORE", 9)}),
        "B": (3, {("ORE", 8)}),
        "C": (5, {("ORE", 7)}),
        "AB": (1, {("A", 3), ("B", 4)}),
        "BC": (1, {("B", 5), ("C", 7)}),
        "CA": (1, {("C", 4), ("A", 1)}),
        "FUEL": (1, {("AB", 2), ("BC", 3), ("CA", 4)}),
    }
    assert ore_required(eg2) == 165
    eg3 = {
        "STKFG": (1, {("VPVL", 2), ("FWMGM", 7), ("CXFTF", 2), ("MNCFX", 11)}),
        "VPVL": (8, {("NVRVD", 17), ("JNWZP", 3)}),
        "FUEL": (
            1,
            {
                ("STKFG", 53),
                ("MNCFX", 6),
                ("VJHF", 46),
                ("HVMC", 81),
                ("CXFTF", 68),
                ("GNMV", 25),
            },
        ),
        "FWMGM": (5, {("VJHF", 22), ("MNCFX", 37)}),
        "NVRVD": (4, {("ORE", 139)}),
        "JNWZP": (7, {("ORE", 144)}),
        "HVMC": (
            3,
            {("MNCFX", 5), ("RFSQX", 7), ("FWMGM", 2), ("VPVL", 2), ("CXFTF", 19)},
        ),
        "GNMV": (6, {("VJHF", 5), ("MNCFX", 7), ("VPVL", 9), ("CXFTF", 37)}),
        "MNCFX": (6, {("ORE", 145)}),
        "CXFTF": (8, {("NVRVD", 1)}),
        "RFSQX": (4, {("VJHF", 1), ("MNCFX", 6)}),
        "VJHF": (6, {("ORE", 176)}),
    }
    assert ore_required(eg3) == 180697

def test_part2() -> None:
    """
    Examples for Part 2.
    """
    eg3 = {
        "STKFG": (1, {("VPVL", 2), ("FWMGM", 7), ("CXFTF", 2), ("MNCFX", 11)}),
        "VPVL": (8, {("NVRVD", 17), ("JNWZP", 3)}),
        "FUEL": (
            1,
            {
                ("STKFG", 53),
                ("MNCFX", 6),
                ("VJHF", 46),
                ("HVMC", 81),
                ("CXFTF", 68),
                ("GNMV", 25),
            },
        ),
        "FWMGM": (5, {("VJHF", 22), ("MNCFX", 37)}),
        "NVRVD": (4, {("ORE", 139)}),
        "JNWZP": (7, {("ORE", 144)}),
        "HVMC": (
            3,
            {("MNCFX", 5), ("RFSQX", 7), ("FWMGM", 2), ("VPVL", 2), ("CXFTF", 19)},
        ),
        "GNMV": (6, {("VJHF", 5), ("MNCFX", 7), ("VPVL", 9), ("CXFTF", 37)}),
        "MNCFX": (6, {("ORE", 145)}),
        "CXFTF": (8, {("NVRVD", 1)}),
        "RFSQX": (4, {("VJHF", 1), ("MNCFX", 6)}),
        "VJHF": (6, {("ORE", 176)}),
    }
    assert fuel_from_trillion_ore(eg3) == 5586022

def main() -> None:
    """
    Calculate and output the solutions based on the real puzzle input.
    """
    data = aocd.get_data(year=2019, day=14)
    recipes = read_recipes(data)

    print(f"Part 1: {ore_required(recipes)}")
    print(f"Part 2: {fuel_from_trillion_ore(recipes)}")

if __name__ == "__main__":
    main()


from fastapi import FastAPI, WebSocket, Request
import uvicorn
import json
from fastapi.templating import Jinja2Templates

app = FastAPI()

templates = Jinja2Templates(directory='templates')

@app.get("/")
async def root(request: Request):
    return templates.TemplateResponse('index.html', {'request': request})

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    id = 1
    while True:
        data_1 = await websocket.receive_json()
        data_2 = {}
        data_2[id] = data_1['id']
        data_2 = json.dumps(data_2)
        await websocket.send_json(data_2)

if __name__ == '__main__':
    uvicorn.run('main:app', port=8000, host='127.0.0.1', reload=True)


#!/usr/bin/python
from sh import grep, avahi_resolve_address, wget
import os, sys, re, sh

OUI_FILE = '/usr/share/arp-scan/ieee-oui.txt'

stdout = sys.stdout.write

oui_url = 'https://raw.githubusercontent.com/royhills/arp-scan/master/ieee-oui.txt'

mac_regex = re.compile('(?:[0-9a-fA-F]:?){12}')
ipv4_like_regex = re.compile(r'[0-9]+(?:\.[0-9]+){3}')
valid_ip_octets = lambda ip: all([0<=int(p)<256 for p in ip.split('.')])

def check_oui_file_and_prompt(oui_file):
    tmp_oui_file = '/tmp/ieee-oui.txt'
    home_oui_file = os.path.expanduser('~') + 'ieee-oui.txt'
    paths = [oui_file, tmp_oui_file, home_oui_file]
    for filepath in paths:
        if os.path.exists(filepath):
            return filepath
    stdout('\nCould not find any of {}. Execute this command to download:\n\t{}\n'
           ''.format(paths, 'wget {} -O ~/ieee-oui.txt'.format(oui_url)))
    return None

def extract_first_ip(text):
    dot_matches = re.findall(ipv4_like_regex, text)
    ips = [x for x in dot_matches if valid_ip_octets(x)]
    return ips[0] if ips else None

def sh_host(ip):
    call = sh.host(ip, _ok_code=[0,1])
    if call.exit_code != 0:
        return None
    else:
        host = str(call).split(' pointer ')[1:]
        return host[0] if host else None

def extract_first_mac(text):
    macs = re.findall(mac_regex, text)
    return macs[0] if macs else None

def get_vendor(mac):
    if not mac:
        return None
    grepped = grep(
        mac.upper().replace(':','')[:6],
        OUI_FILE,
        _ok_code=[0,1]
    )
    return str(grepped)[7:] or '<unknown vendor>'

def dict_to_string(d):
    return '{{{}}}'.format(
        ', '.join(
            ["{}: {}".format(k, v) for k, v in d.iteritems()]
        )
    )

def process_stdin():
    lines = sys.stdin.readlines()
    ips = map(extract_first_ip, lines)
    macs = map(extract_first_mac, lines)
    _avahi_call = [(ip, avahi_resolve_address(ip, _iter=True)) for ip in ips if ip]
    avahi = {k:str(v) for (k,v) in _avahi_call}
    for line in lines:
        ip = extract_first_ip(line)
        mac = extract_first_mac(line)
        half_sep = '\t' if '\t' in line else ' '*4
        sep = half_sep * 2
        extra = ''
        if mac:
            extra += sep + get_vendor(mac)
        if ip:
            avahi_host = avahi[ip].split('\t')[1] if avahi.get(ip, False) else ''
            extra += sep + (sh_host(ip) or avahi_host or '<unknown host>')
        if extra:
            sys.stdout.write((line + half_sep + extra).replace('\n', '') + '\n')
        else:
            sys.stdout.write(line)

help_text = ''\
"""
IP&MAC address lookup script
============================
This script will find the the first IP and MAC address in lines
from stdin, and append MAC-vendor and IP hostname to the lines if
found.

Usage:
------
File:
    cat my_ip_and_mac_addresses.txt | ./ipmac_lookup.py

Local subnet:
    sudo arp-scan 192.168.1.0/24 | ./ipmac_lookup.py

Local ip address:
    arp 192.168.1.1 | ./ipmac_lookup.py
 Output:
    Address          HWtype  HWaddress           Flags Mask  Iface
    192.168.1.1    ether   38:2c:4b:aa:bb:cc   C           eth0    ASUSTek COMPUTER INC.   somehostname.local

The script currently processes all stdin at once, because executing
`avahi-resolve-address` on a bunch of IP addresses is a lot faster than
potentially timing out on one IP address after another.

TODO: Rewrite in order to choose if we want to process stdin one line at a
      time, which would make it possible to use the script as a filter-step
      in a tail -f command, for example.

A prerequisite is a ieee-oui.txt file for performing the MAC <--> vendor
lookup.
The official IEEE link is dead slow. A faster mirror can be found at:
https://raw.githubusercontent.com/royhills/arp-scan/master/ieee-oui.txt
"""

if __name__ == '__main__':
    if os.isatty(0):
        stdout(help_text + '\n')
    else:
        OUI_FILE = check_oui_file_and_prompt(OUI_FILE)
        process_stdin()


#!/usr/bin/env python3
# @Time    : 2020-02-14
# @Author  : caicai
# @File    : request.py
import requests
import urllib3

urllib3.disable_warnings()

from myscan.lib.core.data import cmd_line_options, logger
from myscan.lib.core.common import getredis, gethostportfromurl
from myscan.lib.core.block_info import block_info
from time import sleep
from random import uniform
from myscan.pocs.search import searchmsg
from myscan.config import scan_set
import copy
from myscan.lib.core.const import key_unquote

def request(**kwargs_sour):
    return do_req(None, **kwargs_sour)

def request_session(session, **kwargs_sour):
    return do_req(session, **kwargs_sour)

def do_req(session, **kwargs_sour):
    kwargs = copy.deepcopy(kwargs_sour)
    if kwargs.get("quote", None) is not None:
        if isinstance(kwargs.get("quote"), bool):
            if not kwargs.get("quote"):
                kwargs["url"] = key_unquote + kwargs.get("url", "")
        else:
            logger.warning("requests quote args need bool")
            return
        del kwargs["quote"]
    # print("start:",kwargs)
    if not kwargs.get("verify", None):
        kwargs["verify"] = False
    if not kwargs.get("timeout", None):
        kwargs["timeout"] = 12
    if not kwargs.get("headers", None):
        kwargs["headers"] = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; WOW64; rv:38.0) Gecko/20100101 Firefox/38.0"}
    elif not kwargs.get("headers").get("User-Agent", None):
        kwargs["headers"]["User-Agent"] = "Mozilla/5.0 (Windows NT 10.0; WOW64; rv:38.0) Gecko/20100101 Firefox/38.0"

    if cmd_line_options.proxy:
        kwargs["proxies"] = cmd_line_options.proxy
    if cmd_line_options.timeout:
        kwargs["timeout"] = cmd_line_options.timeout
    # print("end:",kwargs)
    if kwargs.get('data', None):
        if isinstance(kwargs.get("data"), str):
            kwargs["data"] = kwargs["data"].encode("utf-8", "ignore")
    r = None
    h, p = gethostportfromurl(kwargs.get("url"))
    block = block_info(h, p)
    if block.is_block():
        return None
    # retry
    for x in range(cmd_line_options.retry + 1):
        try:
            if session is None:
                r = requests.request(**kwargs)
            else:
                r = session.request(**kwargs)
            break
        except requests.exceptions.ConnectTimeout:
            pass
            # logger.debug("request connect timeout :{}".format(kwargs["url"]))
        except requests.exceptions.ReadTimeout:
            pass
            # logger.debug("request read timeout :{}".format(kwargs["url"]))
        except Exception as ex:
            # print(kwargs)
            logger.debug("Request error url:{} error:{}".format(kwargs["url"], ex))
        sleep(uniform(0, 0.2))
    red = getredis()
    red.hincrby("count_all", "request", amount=1)
    if r != None:
        block.push_result_status(0)
        if scan_set.get("search_open", False):
            s = searchmsg(r)
            s.verify()
            # s.saveresult()
    else:
        block.push_result_status(1)
        red.hincrby("count_all", "request_fail", amount=1)
    return r


import cv2

img = cv2.imread("images.jpeg")

gaussian = cv2.GaussianBlur(img,(15,15),0)#  (yapıcalak gorsel, çekirdek boyutu ,sigma)  genel blurlama yapıyor.

median = cv2.medianBlur(img,5) # gürültüyü gidermede kullanılır.

bila = cv2.bilateralFilter(img,20,120,75) #kenar koruyarak blurlama performansı gaustan iyidr. kenarları korur. gaus her şeyi blurlar.

blur = cv2.blur(img,(5,5))

cv2.imshow("blur",blur)
cv2.imshow("bila",bila)
cv2.imshow("median",median)
cv2.imshow("gassuian",gaussian)
cv2.imshow("image",img)
cv2.waitKey(0)
cv2.destroyAllWindows()


import setuptools

from tensorcircuit import __version__, __author__
from tensorcircuit.utils import is_m1mac

with open("README.md", "r") as fh:
    long_description = fh.read()

install_requires = ["numpy", "scipy", "tensornetwork", "networkx"]

if not is_m1mac():
    install_requires.append("tensorflow")
    # avoid the embarassing macos M1 chip case, where the package is called tensorflow-macos

setuptools.setup(
    name="tensorcircuit",
    version=__version__,
    author=__author__,
    author_email="shixinzhang@tencent.com",
    description="Quantum circuits on top of tensor network",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://github.com/tencent-quantum-lab/tensorcircuit",
    packages=setuptools.find_packages(),
    include_package_data=True,
    install_requires=install_requires,
    tests_require=[
        "pytest",
        "pytest-lazy-fixture",
        "pytest-cov",
        "pytest-benchmark",
        "pytest-xdist",
    ],
    classifiers=[
        "Programming Language :: Python :: 3",
        "Operating System :: OS Independent",
    ],
)


#!/usr/bin/python
# -*- coding:utf-8 -*-
import datetime
import traceback

from actions.applet import find_member_by_open_id
from actions.applet.utils import do_diamond_reward
from db import SOURCE_MEMBER_DIAMOND_DAILY_SHARE_TIMES_LIMIT, SOURCE_MEMBER_DIAMOND_DAILY_SHARE, \
    SOURCE_MEMBER_DIAMOND_DICT
from db.model_utils import get_cache_share_times, set_cache_share_times
from db.models import GameDiamondReward
from logger import log_utils
from web import decorators, WechatAppletHandler

logger = log_utils.get_logging('a_miniapp')

class MiniappShareAwardViewHandler(WechatAppletHandler):
    """ 分享好友奖励会员钻石"""

    @decorators.render_json
    @decorators.wechat_applet_authenticated
    async def post(self):
        r_dict = {'code': 0}
        open_id = self.get_i_argument('open_id', None)
        try:
            if open_id:
                member = await find_member_by_open_id(open_id)
                if member:
                    award_title = ''
                    award_diamonds = ''
                    share_times = await GameDiamondReward.find_one(
                        filtered={'source': SOURCE_MEMBER_DIAMOND_DAILY_SHARE_TIMES_LIMIT})
                    value = get_cache_share_times(member.cid)
                    if value < share_times.quantity:
                        end_datetime = datetime.datetime.now().replace(hour=23, minute=59, second=59, microsecond=999)
                        current_datetime = datetime.datetime.now()
                        timeout = (end_datetime - current_datetime).seconds
                        # 分享奖励
                        _, m_diamond = await do_diamond_reward(member, SOURCE_MEMBER_DIAMOND_DAILY_SHARE)
                        # 设置更新已奖励次数
                        set_cache_share_times(member.cid, timeout)
                        if member.diamond is None:
                            member.diamond = m_diamond
                        else:
                            member.diamond = member.diamond + m_diamond
                        await member.save()
                        award_title = SOURCE_MEMBER_DIAMOND_DICT.get(SOURCE_MEMBER_DIAMOND_DAILY_SHARE)
                        award_diamonds = m_diamond
                    # 用户的钻石
                    r_dict['diamonds_num'] = member.diamond
                    r_dict['award_title'] = award_title
                    r_dict['award_diamonds'] = award_diamonds
                    r_dict['code'] = 1000
                else:
                    r_dict['code'] = 1002
            else:
                r_dict['code'] = 1001
        except Exception:
            logger.error(traceback.format_exc())

        return r_dict


# -*- coding: utf-8 -*-
"""
Created on Tue Feb 16 16:50:52 2016

@author: pchambers
"""

# transonic_airliner.py ========================================================
# Parametric geometry covering some of the design space of 'tube and wing' type
# transonic airliners. Run this as a script to generate geometries approximating
# those of the B787 Dreamliner (-8 and -9) and Airbus A380 (uncomment the appro-
# priate line at the end of this file - see the header of the function for brief
# explanations of the design variables, if you wish to build your own instances
# of this parametric model.
#
# NOTE: For good results in terms of visualisation, select the 'Rendered' view
# in your Rhinoceros viewports. Better still, download the free Neon render
# plugin, which should put a 'Raytraced with Neon' entry in your viewport type
# menu.
# ==============================================================================
# AirCONICS
# Aircraft CONfiguration through Integrated Cross-disciplinary Scripting
# version 0.2.1
# Andras Sobester, 2015.
# Bug reports to a.sobester@soton.ac.uk or @ASobester please.
# ==============================================================================
from OCC.Graphic3d import Graphic3d_NOM_ALUMINIUM
from OCC.gp import gp_Ax1, gp_Pnt, gp_Dir

from airconics.examples.wing_example_transonic_airliner import *
import airconics.examples.boxwing as bw
from airconics.base import AirconicsShape, AirconicsCollection
from airconics.examples.tailplane_example_transonic_airliner import *
from airconics import liftingsurface, engine, fuselage_oml
import airconics.AirCONICStools as act

def transonic_airliner(display=None,
                       Propulsion=1,
                       EngineDia=2.9,
                       FuselageScaling=[55.902, 55.902, 55.902],
                       NoseLengthRatio=0.182,
                       TailLengthRatio=0.293,
                       WingScaleFactor=44.56,
                       WingChordFactor=1.0,
                       Topology=1,
                       SpanStation1=0.35,
                       SpanStation2=0.675,
                       EngineCtrBelowLE=0.3558,
                       EngineCtrFwdOfLE=0.9837,
                       Scarf_deg=3):
    """
    Parameters
    ----------
    Propulsion - int
        1 - twin
        2 - quad
    EngineDia - float
        Diameter of engine intake highlight
    FuselageScaling - list, length 3
        Fx, Fy, Fz scaling factors of fuselage
    NoseLengthRatio - scalar
        Proportion of forward tapering section of the fuselage
    TailLengthRatio - scalar
        Proportion of aft tapering section of the fuselage
    WingScaleFactor - scalar
        Scale Factor of the main wing
    WingChordFactor - scalar
        Chord factor of the main wing
    Topology - int
        Topology = 2 should yield a box wing airliner - use with caution
    SpanStation - float
        Inboard engine at this span station
    SpanStation2 - float
        Outboard engine at this span station (ignored if Propulsion=1)
    EngineCtrBelowLE - float
        Engine below leading edge, normalised by the length of the nacelle -
        range: [0.35,0.5]
    EngineCtrFwdOfLE - float
        Engine forward of leading edge, normalised by the length of the nacelle
        range: [0.85,1.5]
    Scarf_deg - scalar # Engine scarf angle

    Returns
    -------
    airliner - 'Aircraft' class instance
        The collection of aircraft parts

    See also
    --------
    class Aircraft
    """

    try:
        Fus = fuselage_oml.Fuselage(NoseLengthRatio, TailLengthRatio,
                                    Scaling=FuselageScaling,
                                    NoseCoordinates=[0, 0, 0])
    except:
        print("Fuselage fitting failed - stopping.")
        return None

    FuselageHeight = FuselageScaling[2]*0.105
    FuselageLength = FuselageScaling[0]
    FuselageWidth  = FuselageScaling[1]*0.106

    if Fus['OML'] is None:
        print("Failed to fit fuselage surface, stopping.")
        return None

#    FSurf = rs.CopyObject(FuselageOMLSurf)

    # Position of the apex of the wing
    if FuselageHeight < 8.0:
         #787:[9.77,0,-0.307]
        WingApex = [0.1748*FuselageLength,0,-0.0523*FuselageHeight]
    else:
        WingApex = [0.1748*FuselageLength,0,-0.1*FuselageHeight] #787:[9.77,0,-0.307]

    # Set up the wing object, including the list of user-defined functions that
    # describe the spanwise variations of sweep, dihedral, etc.
    if Topology == 1:
        NSeg = 10
        Wing = liftingsurface.LiftingSurface(WingApex,
                                             mySweepAngleFunctionAirliner,
                                             myDihedralFunctionAirliner,
                                             myTwistFunctionAirliner,
                                             myChordFunctionAirliner,
                                             myAirfoilFunctionAirliner,
                                             SegmentNo=NSeg,
                                             ScaleFactor=WingScaleFactor,
                                             ChordFactor=WingChordFactor)
        RootChord = Wing.RootChord
    elif Topology == 2:
        NSeg = 101
        Wing = liftingsurface.LiftingSurface(WingApex,
                                             mySweepAngleFunctionAirliner,
                                             bw.myDihedralFunctionBoxWing,
                                             myTwistFunctionAirliner,
                                             myChordFunctionAirliner,
                                             myAirfoilFunctionAirliner,
                                             SegmentNo=NSeg,
                                             ScaleFactor=WingScaleFactor,
                                             ChordFactor=WingChordFactor)
        RootChord = Wing.RootChord

    if Topology == 1:

        # Add wing to body fairing
        WTBFXCentre = WingApex[0] + RootChord/2.0 + RootChord*0.1297 # 787: 23.8
        if FuselageHeight < 8.0:
            #Note: I made changes to this in OCC Airconics to get a better fit
            # - Paul
            WTBFZ = RootChord*0.009 #787: 0.2
            WTBFheight = 1.8*0.1212*RootChord #787:2.7
            WTBFwidth = 1.08*FuselageWidth
        else:

            WTBFZ = WingApex[2] + 0.005*RootChord
            WTBFheight = 1.8 * 0.09*RootChord
            WTBFwidth = 1.15*FuselageWidth

        WTBFlength = 1.167*RootChord #787:26

        WBF_shape = act.make_ellipsoid([WTBFXCentre, 0, WTBFZ], WTBFlength,
                                        WTBFwidth, WTBFheight)
        WBF = AirconicsShape(components={'WBF': WBF_shape})

#        Trim wing inboard section
        CutCirc = act.make_circle3pt([0,WTBFwidth/4.,-45], [0,WTBFwidth/4.,45], [90,WTBFwidth/4.,0])
        CutCircDisk = act.PlanarSurf(CutCirc)
        Wing['Surface'] = act.TrimShapebyPlane(Wing['Surface'], CutCircDisk)
    elif Topology == 2:
#        Overlapping wing tips
        CutCirc = act.make_circle3pt((0,0,-45), (0,0,45), (90,0,0))
        CutCircDisk = act.PlanarSurf(CutCirc)
        Wing['Surface'] = act.TrimShapebyPlane(Wing['Surface'], CutCircDisk)

#   Engine installation (nacelle and pylon)
    NacelleLength = 1.95*EngineDia
    if Propulsion == 1:
        SpanStations = [SpanStation1]
    elif Propulsion == 2:
        SpanStations = [SpanStation1, SpanStation2]

    engines = []
    for i, SpanStation in enumerate(SpanStations):
        EngineSection, HChord = act.CutSect(Wing['Surface'], SpanStation)
        Chord = HChord.GetObject()
        CEP = Chord.EndPoint()
        Centreloc = [CEP.X()-EngineCtrFwdOfLE*NacelleLength,
                    CEP.Y(),
                    CEP.Z()-EngineCtrBelowLE*NacelleLength]
        eng =  engine.Engine(HChord,
               CentreLocation=Centreloc,
               ScarfAngle=Scarf_deg,
               HighlightRadius=EngineDia/2.0,
               MeanNacelleLength = NacelleLength)

        engines.append(eng)

#    # Script for generating and positioning the fin
#    # Position of the apex of the fin
    P = [0.6524*FuselageLength,0.003,FuselageHeight*0.384]
    #P = [36.47,0.003,2.254]55.902

    if Topology == 1:
        ScaleFactor = WingScaleFactor/2.032 #787:21.93
    elif Topology == 2:
        ScaleFactor = WingScaleFactor/3.5

    SegmentNo = 100
    ChordFactor = 1.01#787:1.01
    Fin = liftingsurface.LiftingSurface(P, mySweepAngleFunctionFin,
                                        myDihedralFunctionFin,
                                        myTwistFunctionFin,
                                        myChordFunctionFin,
                                        myAirfoilFunctionFin,
                                        SegmentNo=SegmentNo,
                                        ChordFactor=ChordFactor,
                                        ScaleFactor=ScaleFactor)

#    Create the rotation axis centered at the apex point in the x direction
    RotAxis = gp_Ax1(gp_Pnt(*P), gp_Dir(1, 0, 0))

    Fin.RotateComponents(RotAxis, 90)

    if Topology == 1:
#        Tailplane
        P = [0.7692*FuselageLength,0.000,FuselageHeight*0.29]
        SegmentNo = 100
        ChordFactor = 1.01
        ScaleFactor = 0.388*WingScaleFactor #787:17.3
        TP = liftingsurface.LiftingSurface(P,
                                           mySweepAngleFunctionTP,
                                           myDihedralFunctionTP,
                                           myTwistFunctionTP,
                                           myChordFunctionTP,
                                           myAirfoilFunctionTP,
                                           SegmentNo=SegmentNo,
                                           ChordFactor=ChordFactor,
                                           ScaleFactor=ScaleFactor)

# TODO: Implement windows
#-----------------------------------------------------------------------
#    Note: some of the code below here may be legacy from Rhino
#    # Windows
#
#    # Cockpit windows:
    CockpitWindowTop = 0.305*FuselageHeight

#    print("Making Fuselage Windows...")

#    solids = Fus.CockpitWindowContours(Height=CockpitWindowTop, Depth=6)
#
##     Cut the windows out:
#    for solid in solids:
#        Fus['OML'] = act.SplitShape(Fus['OML'], solid)
#
#
#    print("Cockpit Windows Done.")

#    (Xmin,Ymin,Zmin,Xmax,Ymax,Zmax) = act.ObjectsExtents([Win1, Win2, Win3, Win4])
    CockpitBulkheadX = NoseLengthRatio*FuselageLength*0.9
#
#    CockpitWallPlane = rs.PlaneFromPoints([CockpitBulkheadX, -15,-15],
#    [CockpitBulkheadX,15,-15],
#    [CockpitBulkheadX,-15,15])
#
#    CockpitWall = rs.AddPlaneSurface(CockpitWallPlane, 30, 30)
#
#    if 'WTBF' in locals():
#        rs.TrimBrep(WTBF, CockpitWall)
#
#    rs.DeleteObject(CockpitWall)
#
#    # Window lines
#    WIN = [1]
#    NOWIN = [0]
#
##     A typical window pattern (including emergency exit windows)
#    WinVec = WIN + 2*NOWIN + 9*WIN + 3*NOWIN + WIN + NOWIN + 24*WIN + 2*NOWIN + WIN + NOWIN + 14*WIN + 2*NOWIN + WIN + 20*WIN + 2*NOWIN + WIN + NOWIN + 20*WIN
#    wires = []
#    if FuselageHeight < 8.0:
##         Single deck
#        WindowLineHeight = 0.3555*FuselageHeight
#        WinX = 0.1157*FuselageLength
#        WindowPitch = 0.609
#        WinInd = -1
#        i = 0
#        while WinX < 0.75*FuselageLength:
#            WinInd = WinInd + 1
#            if WinVec[WinInd] == 1 and WinX > CockpitBulkheadX:
#                i=i+1
#                print("Generating cabin window {}".format(i))
#                WinCenter = [WinX, WindowLineHeight]
#                W_wire = Fus.WindowContour(WinCenter)
#                wires.append(W_wire)
##                display.DisplayShape(W_wire, color='black')
##                win_port, win_stbd = Fus.MakeWindow(*WinCenter)
##                print(type(win_port), type(win_stbd))
##
##                display.DisplayShape(win_port, color='Black')
##                display.DisplayShape(win_stbd, color='Black')
##
##                Note: Need to fix makewindow to return windows WinStbd,
##                # WinPort,
##                Fus.MakeWindow(WinX, WindowLineHeight)
##                act.AssignMaterial(WinStbd,"Plexiglass")
##                act.AssignMaterial(WinPort,"Plexiglass")
#            WinX = WinX + WindowPitch

#    TODO: Fuselage big enough to accommodate two decks
#    else:
#        # Lower deck
#        WindowLineHeight = 0.17*FuselageHeight #0.166
#        WinX = 0.1*FuselageLength #0.112
#        WindowPitch = 0.609
#        WinInd = 0
#        while WinX < 0.757*FuselageLength:
#            WinInd = WinInd + 1
#            if WinVec[WinInd] == 1 and WinX > CockpitBulkheadX:
#                WinStbd, WinPort, FuselageOMLSurf = fuselage_oml.MakeWindow(FuselageOMLSurf, WinX, WindowLineHeight)
#                act.AssignMaterial(WinStbd,"Plexiglass")
#                act.AssignMaterial(WinPort,"Plexiglass")
#            WinX = WinX + WindowPitch
#        # Upper deck
#        WindowLineHeight = 0.49*FuselageHeight
#        WinX = 0.174*FuselageLength #0.184
#        WinInd = 0
#        while WinX < 0.757*FuselageLength:
#            WinInd = WinInd + 1
#            if WinVec[WinInd] == 1 and WinX > CockpitBulkheadX:
#                WinStbd, WinPort, FuselageOMLSurf = fuselage_oml.MakeWindow(FuselageOMLSurf, WinX, WindowLineHeight)
#                act.AssignMaterial(WinStbd,"Plexiglass")
#                act.AssignMaterial(WinPort,"Plexiglass")
#            WinX = WinX + WindowPitch
#
#    print("Cabin windows done")
#-----------------------------------------------------------------------

#     Mirror the geometry as required
    Wing2 = Wing.MirrorComponents(plane='xz')
    try:
        # this try section allows box wing i.e. no tailplane
        TP2 = TP.MirrorComponents(plane='xz')
    except:
        pass

    engines_left = []
    for eng in engines:
        engines_left.append(eng.MirrorComponents(plane='xz'))

#    Build the return assembly (note the three methods of assignment: using
#     the parts keyword, using assignment i.e. airliner[name] = value, and
#     using the airliner.AddPart method):
    airliner = AirconicsCollection(parts={'Wing_right': Wing,
                                          'Wing_left': Wing2,
                                          'Fuselage': Fus,
                                          'Fin': Fin})

    try:
#        boxwing exception (no TP or WBF)
        airliner['WBF'] = WBF
        airliner['Tailplane_right'] = TP
        airliner['Tailplane_left'] = TP2
    except:
        pass

    # Loop over the engines and write all components:
    for i, eng in enumerate(engines):
        name = 'engine_right' + str(i+1)
        airliner.AddPart(eng, name)
    for i, eng in enumerate(engines_left):
        name = 'engine_left' + str(i+1)
        airliner.AddPart(eng, name)

#    To display the entire aircraft in the default colour and material:
#    airliner.Display(display)

#    Set individual component colours and materials:
    from OCC.Quantity import (Quantity_NOC_RED4, Quantity_NOC_WHITE,
                              Quantity_NOC_BLUE4, Quantity_NOC_GRAY)
    from OCC.Graphic3d import Graphic3d_NOM_SHINY_PLASTIC
    red = Quantity_NOC_RED4
    white =  Quantity_NOC_WHITE
    blue = Quantity_NOC_BLUE4
    grey = Quantity_NOC_GRAY
    painted = Graphic3d_NOM_SHINY_PLASTIC   # Gives a painted (ish) appearance

    airliner['Fin'].Display(display, color=red, material=painted)

    try:
#        Boxwing exceptions:
        airliner['Tailplane_left'].Display(display, color=red, material=painted)
        airliner['Tailplane_right'].Display(display, color=red, material=painted)
        airliner['WBF'].Display(display, color=white,  material=painted)
    except:
        pass

    airliner['Wing_left'].Display(display, color=white,  material=painted)
    airliner['Wing_right'].Display(display, color=white,  material=painted)
    airliner['Fuselage'].Display(display, color=white,  material=painted)

    for eng in engines + engines_left:
        display.DisplayShape(eng['Spinner'], color='black')
        display.DisplayShape(eng['Nacelle'], color=blue, material=painted)
        display.DisplayShape(eng['BypassDisk'], color=grey, material=painted)
        display.DisplayShape(eng['FanDisk'], color=grey, material=painted)
        display.DisplayShape(eng['TailCone'], color=grey, material=painted)
        display.DisplayShape(eng['Pylon_symplane'], color=white, material=painted)

    return airliner

if __name__ == "__main__":
    from OCC.Display.SimpleGui import init_display
    display, start_display, add_menu, add_function_to_menu = init_display()

#    A few examples, instances of this parametric aircraft geometry:

#    '787-8'
    # Airliner = transonic_airliner(display)

#    '787-9'
    # Airliner = transonic_airliner(display,
    #                              FuselageScaling=[61.9, 55.902, 55.902])

#    'A380'
    Airliner = transonic_airliner(display, Propulsion=2,
                                  FuselageScaling=[70.4, 67.36, 80.1],
                                  WingScaleFactor=59.26)

#    This, for now, is just intended to stretch the model a bit in terms of
#    topological variety - it is a box wing version of the 787-8. There is
#    no serious design intent here, merely a hint of some of the possibilities
#    in this model.
#    NOTE: this model is rather slow in occ_airconics
#    transonic_airliner(display=display, WingScaleFactor = 66,
#                       WingChordFactor = 0.5, Topology =2)

    # Write step file:
#    import os
#    path = '.'
#    filename = os.path.join(path, 'Airliner.stp')
#    Airliner.Write(filename, single_export=True)    # This function checks filename extension

#     or stl export:
#    filename = os.path.join(path, 'Airliner.stl')
#    Airliner.Write(filename, single_export=True)

    # To output the rendering as pdf:
#    from OCC.Graphic3d import (Graphic3d_EF_PDF,
#                           Graphic3d_EF_PostScript)
#    display.View.View().GetObject().Export('./Airliner.pdf', Graphic3d_EF_PDF)

    start_display()