"""
function for plotting the AbMetaAnalysis results
"""

# Info
__author__ = 'Boaz Frankel'

# Imports
import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.lines import Line2D
import seaborn as sns
import math
from changeo.Gene import getFamily
import typing

def plot_compare_to_reference_cdr3_df(df, figsize=(15, 10)):
    """
    plot the results of compare_to_reference_df
    :param df: output data frame of compare_to_reference_df
    :param figsize: size of the figures
    :return:
    """
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=figsize, sharex=True)

    ax1.bar(range(0, len(df) * 2, 2), df['CASE'].absolute)
    ax1.bar(range(1, len(df) * 2, 2), df['CTRL'].absolute)
    ax1.set_xticks([])
    ax1.set_ylabel('Matches', fontsize=20)
    ax1.legend(labels=['CASE', 'CTRL'], fontsize=20)
    ax1.set_yticks(
        range(df.xs('absolute', level=1, axis=1).max().max() + 1),
        range(df.xs('absolute', level=1, axis=1).max().max() + 1),
        fontsize=20
    )
    ax1.grid(axis='y')
    ax1.set_title('(A)', fontsize=20, loc='left')

    ax2.bar(range(0, len(df) * 2, 2), df['CASE'].support)
    ax2.bar(range(1, len(df) * 2, 2), df['CTRL'].support)
    ax2.set_xticks(np.arange(0.5, len(df) * 2, 2), df.index, fontsize=20)
    ax2.set_ylabel('Support', fontsize=20)
    ax2.set_xlabel('Junction AA length', fontsize=20)
    ax2.set_yticks(
        range(0, math.ceil(df.xs('support', level=1, axis=1).max().max() / 200) * 200, 200),
        range(0, math.ceil(df.xs('support', level=1, axis=1).max().max() / 200) * 200, 200),
        fontsize=20
    )
    ax2.grid(axis='y', ls='--')
    ax2.set_title('(B)', fontsize=20, loc='left')

def boxplot_features(
    by_subj_df,
    by_study_df,
    aliase2study,
    case_aliases,
    ctrl_aliases,
    colors,
    axes,
    case_ctrl_color_diff=3,
    outlier_th=2.5,
    y='prop',
    ylabel="Freq"
):
    """
    2 row figure - upper banners are boxplots of motifs abundance between samples of the same study, lower banners are boxplots of motifs
    of average per study motif abundance.
    :param by_subj_df: by_subj motif analysis df, output of get_motif_analysis
    :param by_study_df: by_study motif analysis df, output of get_motif_analysis
    :param aliase2study: mapping of study id to alias name
    :param case_aliases: list of case studies aliases
    :param ctrl_aliases: list of ctrl studies aliases
    :param colors: colors for the case and ctrl aliases
    :param axes: subplots axes to plot - must match the dimension of teh input data frames
    :param case_ctrl_color_diff: color difference between the case_aliases and ctrl_aliases
    :param outlier_th: zscore threshold to discard outliers from the plotting
    :param y: whe metric to use for the figures
    :param ylabel: string to display as y label
    :return:
    """
    positions = np.arange(0, 10, 1)
    for i, col in enumerate(by_subj_df.loc[:, y].columns):
        ax = axes[0, i]
        for alias_itr, alias in enumerate(case_aliases):
            color = colors[alias_itr]
            ax.boxplot(
                x=by_subj_df.loc[('CASE', aliase2study[alias]), (y, col)].loc[
                    by_subj_df.loc[('CASE', aliase2study[alias]), ('zscore', col)].abs() < outlier_th
                ],
                positions=[positions[alias_itr]],
                boxprops=dict(lw=2, color=color),
                medianprops=dict(lw=2, color='black'),
                whiskerprops=dict(lw=2, color=color),
                capprops=dict(lw=2, color=color),
                widths=0.7,
                flierprops=dict(
                    marker='o', markerfacecolor=color, markersize=2, linestyle='none', markeredgecolor=color
                )
            )
            ax.set_xticks([])
            ax.grid(axis='y', linestyle='--', linewidth=1)
            ax.set_title(col, fontsize=20)
        case_max = alias_itr
        for alias_itr, alias in enumerate(ctrl_aliases):
            color = colors[case_max + 1 + case_ctrl_color_diff + alias_itr]
            ax.boxplot(
                x=by_subj_df.loc[('CTRL', aliase2study[alias]), (y, col)].loc[
                    by_subj_df.loc[('CTRL', aliase2study[alias]), ('zscore', col)].abs() < outlier_th
                ],
                positions=[case_max + 3 + positions[alias_itr]],
                boxprops=dict(lw=2, color=color),
                medianprops=dict(lw=2, color='black'),
                whiskerprops=dict(lw=2, color=color),
                capprops=dict(lw=2, color=color),
                widths=0.7,
                flierprops=dict(
                    marker='o', markerfacecolor=color, markersize=2, linestyle='none', markeredgecolor=color
                )
            )
            ax.set_xticks([])
            ax.grid(axis='y', linestyle='--', linewidth=1)
            ax.set_title(col, fontsize=20)
        ctrl_max = alias_itr
        if i == 0:
            if ylabel:
                ax.set_ylabel('Sample ' + ylabel, fontsize=20)
            ax.tick_params(labelsize=20)

        ax = axes[1, i]
        ax.boxplot(
            x=by_study_df.loc['CASE', (y, col)],
            positions=[max(ctrl_max, case_max) / 2],
            medianprops=dict(color='gray'),
            widths=max(ctrl_max, case_max) - 2,
            showfliers=False
        )
        case_positions = np.arange(max(ctrl_max, case_max) / 2 - case_max / 2 * 0.3, 100, 0.3)
        ax.scatter(
            x=case_positions[:sum(by_study_df.loc['CASE', ('zscore', col)].abs() < outlier_th)],
            y=by_study_df.loc['CASE', (y, col)].loc[
                by_study_df.loc['CASE', ('zscore', col)].abs() < outlier_th
            ],
            color=colors[:case_max + 1]
        )
        ax.boxplot(
            x=by_study_df.loc['CTRL', (y, col)],
            positions=[max(ctrl_max, case_max) + 2 + max(ctrl_max, case_max) / 2],
            medianprops=dict(color='gray'),
            widths=max(ctrl_max, case_max) - 2,
            showfliers=False
        )
        ctrl_positions = np.arange(2 + max(ctrl_max, case_max) * 1.5 - ctrl_max / 2 * 0.3, 100, 0.3)
        ax.scatter(
            x=ctrl_positions[:sum(by_study_df.loc['CTRL', ('zscore', col)].abs() < outlier_th)],
            y=by_study_df.loc['CTRL', (y, col)].loc[
                by_study_df.loc['CTRL', ('zscore', col)].abs() < outlier_th
            ],
            color=colors[case_max + 1 + case_ctrl_color_diff: case_max + 1 + case_ctrl_color_diff + ctrl_max + 1]
        )
        ax.grid(axis='y', linestyle='--', linewidth=1)
        ax.set_xticks(
            [max(ctrl_max, case_max) / 2, max(ctrl_max, case_max) + 2 + max(ctrl_max, case_max) / 2],
            ['CASE', 'CTRL'], fontsize=20, rotation='vertical'
        )
        if i == 0:
            if ylabel:
                ax.set_ylabel('Study ' + ylabel, fontsize=20)
            ax.tick_params(labelsize=20)
    y_max_row1, y_max_row2 = 0, 0
    for col_itr in range(len(by_subj_df.loc[:, y].columns)):
        y_max_row1 = max(y_max_row1, axes[0, col_itr].get_ylim()[1])
        y_max_row2 = max(y_max_row2, axes[1, col_itr].get_ylim()[1])
    for col_itr in range(len(by_subj_df.loc[:, y].columns)):
        axes[0, col_itr].set_ylim(0, y_max_row1)
        axes[1, col_itr].set_ylim(0, y_max_row2)

def plot_results(results: pd.DataFrame):
    """
    create box plot of cross validation results
    :param results:  data frame - output of test_rep_classifier
    :return:
    """
    res_dict = {
        'Accuracy': results.accuracy_score,
        'Precision': results.precision_score,
        'Recall': results.recall_score
    }
    fig, ax = plt.subplots()
    ax.boxplot(res_dict.values(), flierprops=dict(markerfacecolor='r', marker='d'))
    ax.set_xticklabels(res_dict.keys())
    ax.set_yticks(np.arange(0, 1.1, 0.1))
    plt.grid(axis='y', linestyle='--', linewidth=1)

def plot_results_side_by_side(results_l, name_l, figsize=(9, 6)):
    """
    create multiple banners plot of different cross validation results
    :param results_l: list of results data frame - outputs of test_rep_classifier
    :param name_l: list of labels matching each results data frame
    :param figsize: figure size tuple
    :return:
    """
    fig, ax = plt.subplots(figsize=figsize)
    colors = sns.color_palette("Set2", len(results_l))

    bp_l = []
    for i, (results, name) in enumerate(zip(results_l, name_l)):
        res_dict = {
            'Accuracy': results.accuracy_score,
            'Precision': results.precision_score,
            'Recall': results.recall_score
        }
        bp = ax.boxplot(
            res_dict.values(),
            positions=[i + 1, i + 1 + len(results_l) + 1, i + 1 + 2 * len(results_l) + 2],
            boxprops=dict(linewidth=2, color=colors[i]),
            medianprops=dict(linewidth=4, color=colors[i]),
            flierprops=dict(markeredgecolor=colors[i])
        )
        bp_l.append(bp)

    leg = ax.legend(
        [bp['boxes'][0] for bp in bp_l], name_l, handlelength=3, fontsize=20
    )
    for i, legobj in enumerate(leg.legendHandles):
        legobj.set_linewidth(10.0)
        legobj.set_color(colors[i])

    ax.set_xticks([(i + 1 + i * len(results_l) + (len(results_l) - 1) / 2) for i in range(3)])
    ax.set_xticklabels(['Accuracy', 'Precision', 'Recall'], fontsize=20, fontweight='bold')
    ax.set_yticks(np.arange(0, 1.1, 0.1))
    ax.set_yticklabels(np.round(np.arange(0, 1.1, 0.1), 2), fontsize=20)
    plt.grid(True, axis='y', linestyle='--', linewidth=1)

def plot_features_side_by_side(feature_table_dict: dict, labels: pd.Series):
    """
    create multiple banners plot, each banner is a scatter plot of the between-samples clusters public case index vs public ctrl index
    :param feature_table_dict: dictionary of public clusters feature tables - outputs of build_feature_table
    :param labels: labels of the samples
    :return:
    """
    fig, ax_arr = plt.subplots(1, 2, figsize=(20, 10))
    all_points = pd.DataFrame()
    max_tick = 0
    for i, (name, feature_table) in enumerate(feature_table_dict.items()):

        ax = ax_arr[i]
        points = pd.value_counts(
            [(x[1].loc[labels.index[labels]].sum(), x[1].loc[labels.index[~labels]].sum()) for x in feature_table.items()]
        )
        if (0, 0) in points.index:
            points = points.drop(index=(0, 0))
        s = points.apply(lambda x: math.log(x, 2) * 100 + 100)
        df = pd.DataFrame(
            [list(map(lambda x: x[0], points.index)), list(map(lambda x: x[1], points.index)), s.to_list()],
            index=['x', 'y', 's']
        ).transpose()
        all_points = pd.concat([all_points, df])
        ax.scatter(
            x='x', y='y', s='s', data=df.loc[(df.y > 0) & (df.x > 0)]
        )
        ax.scatter(
            x='x', y='y', s='s', data=df.loc[df.x == 0]
        )
        ax.scatter(
            x='x', y='y', s='s', data=df.loc[df.y == 0]
        )
        max_tick = max(max_tick, int(max(ax.get_xticks().max(), ax.get_yticks().max())))
        ax.set_xticks(range(max_tick), range(max_tick), fontsize=20)
        ax.set_yticks(range(max_tick), range(max_tick), fontsize=20)
        ax.set_aspect('equal')
        ax.set_xlabel("")
        ax.set_ylabel("")
        ax.set_title('(' + name + ')', fontsize=20, loc='left')

    for ax in ax_arr:
        ax.set_xlim(-1, max_tick)
        ax.set_ylim(-1, max_tick)

        ax.set_xticks(range(max_tick))
        ax.set_xticklabels(range(max_tick), fontsize='20')
        ax.set_yticks(range(max_tick))
        ax.set_yticklabels(range(max_tick), fontsize='20')
        ax.set_aspect('equal')
        ax.grid(axis='both', ls='--')

    fig.supxlabel('Case', fontsize=20, y=0.08)
    fig.supylabel('Control', fontsize=20, x=0.08)

    c = ax.scatter(x='x', y='y', s='s', data=all_points, edgecolors='w', c='w', alpha=0)
    elements = c.legend_elements('sizes', num=10, func=lambda x: ((x - 100) / 100), color='gray', alpha=1)
    legend = fig.legend(
        *(np.flip(elements[0]), np.flip(elements[1])),
        labelspacing=1.5,
        handletextpad=1,
        title=r'$Log_{2}$',
        title_fontsize=20,
        fontsize=15
    )
    legend.get_frame().set_linewidth(0.0)

def create_studies_legend_handles(aliases: list, colors):
    """
    auxiliary method to create handlers for legend of studies aliases
    :param aliases: list of study id aliases
    :param colors: list of colors matching the aliases
    :return:
    """
    handles = []
    for alias_itr, alias in enumerate(aliases):
        handles.append(
            Line2D(
                [0], [0], marker='o', color='w', label=alias, markerfacecolor=colors[alias_itr], markersize=10
            )
        )
    return handles

def boxplot_junction_aa_length(
    axes,
    case_airr_seq_df: pd.DataFrame,
    ctrl_airr_seq_df: pd.DataFrame,
    v_families: typing.Iterable,
    case_aliases: typing.Iterable,
    ctrl_aliases: typing.Iterable,
    aliase2study: dict,
    colors: list,
    case_ctrl_color_diff=3
):
    """
    same as boxplot_features but ad-hoc implementation for junction AA length
    :param axes: array of matplot lib axes
    :param case_airr_seq_df: airr-seq data frame of case sequences
    :param ctrl_airr_seq_df: airr-seq data frane if ctrl sequences
    :param v_families: list of v families for which create banners
    :param case_aliases: list of case study ids aliases
    :param ctrl_aliases: list of ctrl study ids aliases
    :param aliase2study: mapping of alias to study id
    :param colors: colors for the study id aliases
    :param case_ctrl_color_diff: color diff between the case and ctrl study id aliases
    :return:
    """
    by_v_family_by_study_junction_len = pd.Series(
        dtype=object, name='v_family', index=pd.MultiIndex.from_product([['CASE', 'CTRL'], v_families])
    )
    for v_family, frame_index in case_airr_seq_df.index.groupby(case_airr_seq_df.v_call_original.apply(getFamily)).items():
        if v_family not in v_families:
            continue
        by_v_family_by_study_junction_len[
            ('CASE', v_family)
        ] = case_airr_seq_df.loc[frame_index].groupby('study_id').apply(lambda x: x.junction_aa_length.astype(int).tolist())
    for v_family, frame_index in ctrl_airr_seq_df.index.groupby(ctrl_airr_seq_df.v_call_original.apply(getFamily)).items():
        if v_family not in v_families:
            continue
        by_v_family_by_study_junction_len[
            ('CTRL', v_family)
        ] = ctrl_airr_seq_df.loc[frame_index].groupby('study_id').apply(lambda x: x.junction_aa_length.astype(int).tolist())
    positions = np.arange(0, 10, 1)
    for i, v_family in enumerate(v_families):
        ax = axes[0, i]
        for alias_itr, alias in enumerate(case_aliases):
            color = colors[alias_itr]
            ax.boxplot(
                x=by_v_family_by_study_junction_len.loc[('CASE', v_family)].loc[aliase2study[alias]],
                positions=[positions[alias_itr]],
                boxprops=dict(lw=2, color=color),
                medianprops=dict(lw=2, color='black'),
                whiskerprops=dict(lw=2, color=color),
                capprops=dict(lw=2, color=color),
                widths=0.7,
                flierprops=dict(
                    marker='o', markerfacecolor=color, markersize=2, linestyle='none', markeredgecolor=color
                )
            )
            ax.set_xticks([])
            ax.grid(axis='y', linestyle='--', linewidth=1)
            ax.set_title(v_family, fontsize=20)
        case_max = alias_itr
        for alias_itr, alias in enumerate(ctrl_aliases):
            color = colors[case_max + 1 + case_ctrl_color_diff + alias_itr]
            ax.boxplot(
                x=by_v_family_by_study_junction_len.loc[('CTRL', v_family)].loc[aliase2study[alias]],
                positions=[case_max + 3 + positions[alias_itr]],
                boxprops=dict(lw=2, color=color),
                medianprops=dict(lw=2, color='black'),
                whiskerprops=dict(lw=2, color=color),
                capprops=dict(lw=2, color=color),
                widths=0.7,
                flierprops=dict(
                    marker='o', markerfacecolor=color, markersize=2, linestyle='none', markeredgecolor=color
                )
            )
            ax.set_xticks([])
            ax.grid(axis='y', linestyle='--', linewidth=1)
            ax.set_title(v_family, fontsize=20)
        ctrl_max = alias_itr
        if i == 0:
            ax.set_ylabel('Study Junction Length', fontsize=20)
            ax.tick_params(labelsize=20)

        ax = axes[1, i]
        ax.boxplot(
            x=sum(by_v_family_by_study_junction_len.loc[('CASE', v_family)].tolist(), []),
            positions=[max(ctrl_max, case_max) / 2],
            medianprops=dict(color='gray'),
            widths=max(ctrl_max, case_max) - 2,
            showfliers=True
        )
        ax.boxplot(
            x=sum(by_v_family_by_study_junction_len.loc[('CTRL', v_family)].tolist(), []),
            positions=[max(ctrl_max, case_max) + 2 + max(ctrl_max, case_max) / 2],
            medianprops=dict(color='gray'),
            widths=max(ctrl_max, case_max) - 2,
            showfliers=True
        )
        ax.grid(axis='y', linestyle='--', linewidth=1)
        ax.set_xticks(
            [max(ctrl_max, case_max) / 2, max(ctrl_max, case_max) + 2 + max(ctrl_max, case_max) / 2],
            ['CASE', 'CTRL'], fontsize=20, rotation='vertical'
        )
        if i == 0:
            ax.set_ylabel('Junction Length', fontsize=20)
            ax.tick_params(labelsize=20)


"""Create reads

Revision ID: 49b4cc7895ba
Revises: 3f30d324d67d
Create Date: 2023-09-20 21:01:20.431436

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = '49b4cc7895ba'
down_revision = '3f30d324d67d'
branch_labels = None
depends_on = None

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('reads',
    sa.Column('id', sa.Integer(), autoincrement=True, nullable=False),
    sa.Column('readable', sa.Integer(), nullable=True),
    sa.Column('readable_id', sa.Integer(), nullable=True),
    sa.Column('user_id', sa.Integer(), nullable=True),
    sa.Column('read_at', sa.DateTime(), nullable=True),
    sa.PrimaryKeyConstraint('id')
    )
    with op.batch_alter_table('reads', schema=None) as batch_op:
        batch_op.create_index(batch_op.f('ix_reads_readable'), ['readable'], unique=False)

    # ### end Alembic commands ###

def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('reads', schema=None) as batch_op:
        batch_op.drop_index(batch_op.f('ix_reads_readable'))

    op.drop_table('reads')
    # ### end Alembic commands ###


from rest_framework import serializers
from rest_framework.relations import SlugRelatedField
from rest_framework.validators import UniqueTogetherValidator

from posts.models import Comment, Post, Group, Follow, User

class CommentSerializer(serializers.ModelSerializer):
    """Класс для преобразования сложных данных в простые типы данных Python,
    которые конвертируются в JSON.
    """
    author = SlugRelatedField(read_only=True, slug_field='username')

    class Meta:
        fields = '__all__'
        model = Comment
        read_only_fields = ('post',)

class FollowSerializer(serializers.ModelSerializer):
    user = SlugRelatedField(slug_field='username',
                            queryset=User.objects.all(),
                            default=serializers.CurrentUserDefault())
    following = serializers.SlugRelatedField(
        slug_field="username", queryset=User.objects.all()
    )

    class Meta:
        fields = ("id", "user", "following")
        model = Follow

        validators = [
            UniqueTogetherValidator(
                queryset=Follow.objects.all(),
                fields=['user', 'following']
            )
        ]

    def validate(self, data):
        """
        Проверяет: 'user не равен 'following'.
        """
        if data['user'] == data['following']:
            raise serializers.ValidationError('Вы не можете'
                                              'подписаться на себя.')
        return data

class PostSerializer(serializers.ModelSerializer):
    """Класс для преобразования сложных данных в простые типы данных Python,
    которые конвертируются в JSON.
    """
    author = SlugRelatedField(slug_field='username',
                              read_only=True)

    class Meta:
        fields = ('id', 'text', 'author',
                  'group', 'image', 'pub_date')
        model = Post

class GroupSerializer(serializers.ModelSerializer):
    """Класс для преобразования сложных данных в простые типы данных Python,
    которые конвертируются в JSON.
    """
    class Meta:
        fields = ('id', 'title', 'slug', 'description')
        model = Group


import mpmath;

class Scorer:

	def score_edge(self,Gain,Sigma):

		G=Gain[0];
		S=Sigma[0];

		if G<=0:
			return mpmath.ln(1);

		if G<S:
			S=G;

		eps = mpmath.ln( 1.0/mpmath.ln(2)) / mpmath.ln(2);

		t0 = mpmath.power(2,-G-eps);

		t1 = mpmath.power(2,-S-eps);
		#print 't1: ',t1;
		t2 = (1.0/G) * mpmath.power(2,-G);
		#print 't2: ',t2;
		t3 = mpmath.power(t1,1) + mpmath.power(t2,1);
		#print 't3: ',t3;
		t4 = -(S)#mpmath.ln(mpmath.power(2,-S));
		#t4 = mpmath.ln(t0 * t3);

		return t4;

	def score_edge_crude(self,Gain,Sigma):
		G=Gain[0];
		S=Sigma[0];
		if G<=0:
			return mpmath.ln(1);

		t0 = mpmath.power(2,-(G+S));
		t1 = mpmath.ln(t0);
		return t1;


import pytest
from unittest.mock import MagicMock, mock_open
from collections import namedtuple

@pytest.fixture()
def mock_object(monkeypatch):
    def create_mock_object(object_path, **kwargs):
        mo = MagicMock(**kwargs)
        monkeypatch.setattr(object_path, mo)
        return mo

    return create_mock_object

@pytest.fixture()
def open_mock(monkeypatch):
    open_mock = mock_open()
    monkeypatch.setattr("builtins.open", open_mock)
    return open_mock

@pytest.fixture()
def os_path_join_file_mock(mock_object):
    os = mock_object("os.path.join")
    os.side_effect = lambda *x: "/".join(x)
    return os

@pytest.fixture()
def json_load_mock(mock_object):
    return mock_object("json.load")

@pytest.fixture()
def argparse_mock(mock_object):
    ap = mock_object("argparse.ArgumentParser")
    ap.return_value = ap
    return ap

@pytest.fixture
def configuration_file():
    def func(headers, sects, errors):
        ConfigurationFile = namedtuple(
            "ConfigurationFile", ["headers", "sects", "errors"]
        )
        return ConfigurationFile(headers, sects, errors)

    return func

@pytest.fixture(scope="session")
def build_parsed_output():
    def parsed_output(output_elements):
        """
        builds expected parser output from provided dict
        :param output_elements: dictionary with {stanza: {option: value, ...}, ...}
        :return: parsed_output
        """
        parsed_output = {}
        for stanza, stanza_value in output_elements.items():
            fake_section = MagicMock()
            fake_section.options = {}
            fake_section.name = stanza
            parsed_output.update({stanza: fake_section})
            for option, value in stanza_value.items():
                fake_setting = MagicMock()
                fake_setting.name = option
                fake_setting.value = value
                parsed_output[stanza].options.update({option: fake_setting})
        return parsed_output

    return parsed_output


import plotly.graph_objects as go

fig = go.Figure(data=[
    go.Mesh3d(
        # Order of vertex matters! in this case is: A-B-C-D-DRONE
        x=[3, 9, 1, 3, 0],
        y=[9, 3, 3, 1, 0],
        z=[0, 0, 0, 0, 6],
        colorbar_title='z',
        # i, j and k give the vertices of triangles
        # here we represent the 4 triangles of the tetrahedron surface
        i=[0, 0, 0, 1, 2, 0],
        j=[2, 1, 1, 3, 3, 2],
        k=[3, 3, 4, 4, 4, 4],
        name='y',
        showscale=True,
        opacity=0.40,
        color="#FF0000"
    )
])

fig.show()


import boltons.fileutils as bf
import yaml
import pkg_resources
from distutils.dir_util import copy_tree
import pprint
from pathlib import Path
import shutil
import re

from jinja2 import Environment, PackageLoader, select_autoescape

resource_package = __name__

env = Environment(
    loader=PackageLoader("docker_app_generator", "templates"),
    autoescape=select_autoescape(["html", "xml"]),
)

def regex_match(string, p):
    return bool(re.search(p,string))

def regex_split(string, p):
    return string[:string.rfind(p)]

env.filters['regex_match'] = regex_match
env.filters['regex_split'] = regex_split

pp = pprint.PrettyPrinter(indent=4)

def load_config_file(filename):
    if is_valid_file(pkg_resources.resource_filename(resource_package, filename)):
        return yaml.load(pkg_resources.resource_string(resource_package, filename))
    else:
        return []

def create_app_structure(path_app):
    bf.mkdir_p(path_app)
    bf.mkdir_p(path_app + "/dockerfiles")

def cleanup_directory(dir):
    if is_valid_directory(dir):
        shutil.rmtree(dir)

def copy_files_old(from_path, to_path):
    copy_tree(from_path, to_path) if is_valid_file(from_path) else None

def copy_files(from_path_base, stack_element, name, to_path_base):
    # copy files
    for file_type in ["config", "files"]:
        from_path = from_path_base + "/" + stack_element + "/" + name + "/" + file_type
        to_path = to_path_base + "/" + file_type + "/" + name
        print("from {}, to {}".format(from_path, to_path))
        copy_tree(from_path, to_path,) if is_valid_file(from_path) else None

    # copy config
    # copy_tree(from_path, to_path) if is_valid_file(from_path) else None

def generate_app_dockerfile(
    path_app, docker_stack, docker_base, docker_flavors, variant_name=None
):
    # print('Dockerfile config {}\ndocker_base: {}\ndocker_stack: {}\n'.format(
    # path_app, docker_base, docker_stack))

    data = {
        "docker_base": docker_base,
        "docker_stack": docker_stack,
        "docker_flavors": docker_flavors,
    }
    generate_file(
        data,
        "Dockerfile",
        path_app
        + "/dockerfiles/"
        + (variant_name + "." if variant_name else "")
        + "Dockerfile",
    )

def generate_build_shell(path_app, docker_image_name, docker_bases, settings):
    # print('Shell config\npath_app: {}\docker_image_name: {}\docker_bases: {}\n'.format(
    # path_app, docker_image_name, docker_bases))
    data = {
        "docker_image_name": docker_image_name,
        "docker_bases": docker_bases,
        "settings": settings,
    }
    generate_file(data, "build.sh", path_app + "/build.sh")

def generate_docker_compose(path_app, app_config, settings):
    data = {
        "app_config": app_config,
        "compose": app_config["compose"],
        "settings": settings,
    }
    generate_file(data, "docker-compose.yml", path_app + "/docker-compose.yml")
    for env in app_config["compose"]["env"]:
        data['env'] = env
        data['compose_env'] = app_config["compose"]["env"][env]
        generate_file(
            data,
            "docker-compose.env.yml",
            path_app + "/docker-compose." + env + ".yml",
        )
    generate_file(
        data, "docker-compose.common.yml", path_app + "/docker-compose.common.yml"
    )

def generate_file(data, filename, target):
    template = env.get_template(filename + ".j2")
    template.stream(data).dump(target)

def is_valid_file(filename):
    p = Path(filename)
    return p.exists()

def is_valid_directory(filename):
    p = Path(filename)
    return p.exists() and p.is_dir()


"""Defines core consumer functionality"""
import logging

import confluent_kafka
from confluent_kafka import Consumer
from confluent_kafka.avro import AvroConsumer
from confluent_kafka.avro.serializer import SerializerError
from tornado import gen

logger = logging.getLogger(__name__)

class KafkaConsumer:
    """Defines the base kafka consumer class"""

    def __init__(
        self,
        topic_name_pattern,
        message_handler,
        is_avro=True,
        offset_earliest=False,
        sleep_secs=1.0,
        consume_timeout=0.1,
    ):
        """Creates a consumer object for asynchronous use"""
        logger.info(f"Creating kafka consumer for {topic_name_pattern}")
        self.topic_name_pattern = topic_name_pattern
        self.message_handler = message_handler
        self.sleep_secs = sleep_secs
        self.consume_timeout = consume_timeout
        self.offset_earliest = offset_earliest

        #
        #
        # TODO: Configure the broker properties below. Make sure to reference the project README
        # and use the Host URL for Kafka and Schema Registry!
        #
        #
        self.broker_properties = {
                #
                # TODO
                #
            'bootstrap.servers': 'PLAINTEXT://localhost:9092',
            "default.topic.config": {"auto.offset.reset": "earliest"},
            'group.id': f'{topic_name_pattern}'
        }

        # TODO: Create the Consumer, using the appropriate type.
        if is_avro is True:
            self.broker_properties["schema.registry.url"] = "http://localhost:8001"
            self.consumer = AvroConsumer(self.broker_properties)
        else:
            self.consumer = Consumer(self.broker_properties)

        # TODO: Configure the AvroConsumer and subscribe to the topics. Make sure to think about
        # how the `on_assign` callback should be invoked.
        logger.info(f"Attempting to subscribe to topic {self.topic_name_pattern}")
        self.consumer.subscribe([self.topic_name_pattern], on_assign=self.on_assign)
        logger.info(f"Successful subscription to topic {self.topic_name_pattern}")

    def on_assign(self, consumer, partitions):
        """Callback for when topic assignment takes place"""
        # TODO: If the topic is configured to use `offset_earliest` set the partition offset to
        # the beginning or earliest
        logger.info(f"on_assign {self.topic_name_pattern} consumer {consumer} and partitions {partitions}")
        for partition in partitions:
            # TODO
            logger.info(f"setting offset ...")
            if self.offset_earliest is True:
#                logger.info("xxx")
                #partition.offset = OFFSET_BEGINNING
                partition.offset = 0 # hanging when using OFFSET_BEGINNING
#                logger.info("yyy")
            logger.info(f"done setting offset")

        logger.info("partitions assigned for %s", self.topic_name_pattern)
        consumer.assign(partitions)

    async def consume(self):
        """Asynchronously consumes data from kafka topic"""
        logger.info(f"Consuming data {self.topic_name_pattern}")
        while True:
            num_results = 1
            while num_results > 0:
                num_results = self._consume()
            logger.info("Sleeping ...")
            await gen.sleep(self.sleep_secs)

    def _consume(self):
        """Polls for a message. Returns 1 if a message was received, 0 otherwise"""
        # TODO: Poll Kafka for messages. Make sure to handle any errors or exceptions.

        # Additionally, make sure you return 1 when a message is processed, and 0 when no message
        # is retrieved.

        logger.info(f"poll message from {self.topic_name_pattern}")
        try:
            message = self.consumer.poll(self.consume_timeout)
            logger.info(f"message returned from poll call is {message}")
        except error:
            logger.info(f'Could not poll message from {self.topic_name_pattern}')
            return 0

        if message is None:
            logger.info(f"No message received")
            return 0
        elif message.error() is not None:
            logger.info(f"Message error: {message.error()}")
            return 0
        else:
            logger.info(f"consumed message {message.topic()}")
            self.message_handler(message)
        return 1

    def close(self):
        """Cleans up any open kafka consumers"""
        # TODO: Cleanup the kafka consumer
        logger.info(f"closing {self.topic_name_pattern}")
        self.consumer.close()


import json
import time

def genere_solution():
    # a*x+b*y=c [10]

    resultat = []
    for a in range(0, 10):
        for b in range(0, 10):
            for c in range(0, 10):

                tmp = {'a': a, 'b': b, 'c': c, 'result': []}
                resultat.append(tmp)
                for x in range(0, 10):
                    for y in range(0, 10):
                        if (a * x + b * y - c) % 10 == 0:
                            print(f'{a}*{x}+{b}*{y}={c}[10]')
                            tmp['result'].append({'x': x, 'y': y})

    return resultat

def main():
    with open('files/solution.json', "w") as f:
        start = time.time()

        res = genere_solution()

        end = time.time()
        elapsed = end * 1000 - start * 1000

        f.write(json.dumps(res, sort_keys=True, indent=4) + "\n")

        print(f'Temps d\'exécution : {elapsed}ms')

if __name__ == '__main__':
    main()


from collections import defaultdict

import numpy as np

from abaqus_python_interface import ABQInterface

from fat_eval.weakest_link.FEM_functions.elements import element_types
from fat_eval.weakest_link.hazard_functions import weibull
from fat_eval.fatigue_materials import materials
from fat_eval.utilities.steel_data import abaqus_fields, SteelData

def setup_weakest_link_evaluator(odb_file, heat_treatment, element_set, instance_name,
                                 symmetry_factor, abaqus):
    abq = ABQInterface(abaqus)
    element_data = abq.get_element_data(odb_file, element_set, instance_name)
    heat_treatment_data = {}
    element_labels = None
    for i, heat_treatment_field in enumerate(abaqus_fields):
        field, _, element_labels = abq.read_data_from_odb(
            odb_file_name=heat_treatment.odb_file_name,
            field_id=heat_treatment_field,
            step_name=heat_treatment.step_name,
            frame_number=heat_treatment.frame_number,
            set_name=heat_treatment.element_set,
            instance_name=heat_treatment.instance,
            get_position_numbers=True
        )
        heat_treatment_data[heat_treatment_field] = field
    elements = {}
    for element_type, element_coordinates in element_data.items():
        for element_label, coords in element_coordinates.items():
            elements[element_label] = element_types[element_type](coords)

    evaluator = WeakestLinkEvaluator(elements, element_labels, SteelData(heat_treatment_data),
                                     symmetry_factor)
    return evaluator

class WeakestLinkEvaluator:
    def __init__(self, elements, element_labels, steel_data, symmetry_factor=1.):
        self.elements = elements
        self.element_labels = element_labels
        self.symmetry_factor = symmetry_factor
        self.steel_data = steel_data

    def evaluate(self, stress_state, material_name, hazard_function=weibull, cycles=2e6):
        material = materials[material_name]
        functional_values = hazard_function(stress_state, self.steel_data, material, cycles=cycles)
        values_dict = defaultdict(list)
        for s, label in zip(functional_values, self.element_labels):
            values_dict[label].append(s)
        integral = 0
        for e_label, values in values_dict.items():
            element = self.elements[e_label]
            for vol, w, val in zip(element.gauss_point_volumes, element.gauss_weights, values):
                integral += val*w*vol

        pf = 1 - np.exp(-integral)
        return 1 - (1 - pf)**self.symmetry_factor


from typing import Dict, List, Optional, Tuple, Callable
import numpy as np

from tensorflow.keras.preprocessing.image import ImageDataGenerator, Iterator

def _sanitize_class_balance(
    classes: List[int], class_balance: Optional[Dict[int, float]] = None
) -> Dict[int, float]:
    if class_balance is None:
        output = {c: 1.0 / len(classes) for c in classes}
    else:
        # normalize the class fractions
        total = np.sum(list(class_balance.values()))
        output = {c: v / total for c, v in class_balance.items()}
    assert np.sum(list(output.values())) == 1.0
    return output

def sample_balanced(
    input_labels: np.ndarray,
    required_samples: int,
    class_balance: Optional[
        Dict[int, float]
    ] = None,  # by default sample classes equally
    shuffle: bool = True,
) -> Dict[int, List[int]]:
    assert input_labels.ndim == 1
    # split input_labels
    classes, indices = np.unique(input_labels, return_index=True)
    classes = list(map(int, classes))
    class_balance = _sanitize_class_balance(
        classes=classes, class_balance=class_balance
    )
    index_dict = {
        c: np.where(input_labels == c)[0]
        if not shuffle
        else np.random.permutation(np.where(input_labels == c)[0])
        for c in classes
    }
    indices_per_class = {
        c: index_dict[c][: int(class_balance[c] * required_samples)] for c in classes
    }
    assert np.sum(list(map(len, indices_per_class.values()))) == required_samples
    return indices_per_class

def sample_balanced_batches(
    input_labels: np.ndarray,
    required_batches: int,
    batch_size: int,
    class_balance: Optional[Dict[int, float]] = None,
    shuffle: bool = True,
) -> np.ndarray:
    assert input_labels.ndim == 1

    indices_per_class = sample_balanced(
        input_labels=input_labels,
        required_samples=required_batches * batch_size,
        class_balance=class_balance,
        shuffle=shuffle,
    )
    classes = list(indices_per_class.keys())

    samples_per_batch = {c: int(class_balance[c] * batch_size) for c in classes}
    assert np.sum(list(samples_per_batch.values())) == batch_size

    sampled = []
    for i in range(required_batches):
        batch = []
        for c in classes:
            samples = samples_per_batch[c]
            batch.append(indices_per_class[c][samples * i : samples * (i + 1)])
        sampled.append(np.concatenate(batch))

    return np.array(sampled)

class UndersamplingIterator(Iterator):
    def __init__(
        self,
        inputs: np.ndarray,
        labels: np.ndarray,
        batch_size: int = 32,
        class_balance: Optional[Dict[int, float]] = None,
        shuffle: bool = True,
        preprocess_fn: Callable = None,
        seed: np.random.RandomState = None,
    ):
        self._inputs = inputs
        self._labels = labels
        self._preprocess_fn = preprocess_fn
        self._labels_argmax = np.argmax(self._labels, axis=1)
        self._batch_size = batch_size
        self._class_balance = class_balance
        self._shuffle = shuffle
        self._req_batches = self._compute_required_batches()
        self._seed = seed
        super(UndersamplingIterator, self).__init__(
            n=self._req_batches,
            batch_size=self._batch_size,
            shuffle=self._shuffle,
            seed=self._seed,
        )
        self._sampled_batch_indices = []
        self._apply_balance_strategy()

    def _compute_required_batches(self) -> int:
        classes, counts = np.unique(self._labels_argmax, return_counts=True)
        self._class_balance = _sanitize_class_balance(
            classes=classes, class_balance=self._class_balance
        )
        min_class = np.argmin(
            [counts[i] / self._class_balance[c] for i, c in enumerate(classes)]
        )
        required_batches = int(
            np.floor(
                int(counts[min_class] / self._class_balance[min_class])
                // self._batch_size
            )
        )
        assert required_batches > 0
        return required_batches

    def on_epoch_end(self):
        """Method called at the end of every epoch.
        """
        self._apply_balance_strategy()

    def _apply_balance_strategy(self):
        self._sampled_batch_indices = sample_balanced_batches(
            input_labels=self._labels_argmax,
            required_batches=self._req_batches,
            batch_size=self._batch_size,
            class_balance=self._class_balance,
            shuffle=self._shuffle,
        )

    def __getitem__(self, index: int) -> Tuple[np.ndarray, np.ndarray]:
        """Gets batch at position `index`.
        Args:
            index: position of the batch in the Sequence.
        Returns:
            A batch
        """
        indices = self._sampled_batch_indices[index]
        X, y = self._inputs[indices, :], self._labels[indices, :]
        if self._preprocess_fn is not None:
            X = self._preprocess_fn(X)
        return X, y

    def __len__(self) -> int:
        return self._req_batches


import numpy as np
import santas_path
import random
from datetime import datetime

def create_particle(particle_size, x_min=0, x_max=4):
    """
    Create a single particle
    :param particle_size:
    :param x_min:
    :param x_max:
    :return:
    """
    particle = x_min + (x_max - x_min) * np.random.uniform(low=0, high=1, size=particle_size)
    return particle

def create_particle_velocities(particle_size, v_min=-4, v_max=4):
    """
    Create the velocities associated with a particle
    :param particle_size:
    :param v_min:
    :param v_max:
    :return:
    """
    velocities = v_min + (v_max - v_min) * np.random.uniform(low=0, high=1, size=particle_size)
    return velocities

def create_pop_particles(pop_size, particle_length):
    """
    Create a population of particles
    :param pop_size:
    :param particle_length:
    :return:
    """
    pop = [create_particle(particle_length) for i in range(0, pop_size)]
    return np.array(pop)

def create_pop_velocities(pop_size, particle_length):
    """
    Create the velocities associated with the particles of a population
    :param pop_size:
    :param particle_length:
    :return:
    """
    pop_vel = [create_particle_velocities(particle_length) for i in range(0, pop_size)]
    return np.array(pop_vel)

def generate_santas_path_from_particles_pop(population):
    """
    Given a population of particles, create plausible paths using the SVP rule
    :param population:
    :return:
    """
    path_populations = np.argsort(population)
    return path_populations

def generate_santas_path_from_particles_pop_complete_dataset(population):
    """

    :param population:
    :return:
    """
    return generate_santas_path_from_particles_pop(population) + 1

def evaluate_paths(paths, cities, not_primes_bool):
    """
    Compute the euclidean distance with penalties value associated with each given path
    :param paths:
    :param cities:
    :param not_primes_bool:
    :return:
    """
    distances = []
    for path in paths:
        distances.append(santas_path.edp(path, cities, not_primes_bool))
    return np.array(distances)

def swap_mutation(pop):
    """
    Exchange the position of two elements
    :param pop:
    :return:
    """
    pos_to_swap = np.random.choice(range(0, len(pop)), size=2, replace=False)
    pop[pos_to_swap[0]], pop[pos_to_swap[1]] = pop[pos_to_swap[1]], pop[pos_to_swap[0]]
    return pop

def mutate_pop_elements(pop, ro):
    """
    Apply swap mutation to 10% of population indivuals
    :param pop:
    :param ro:
    :return:
    """
    elements_to_mutate = np.random.choice(range(0, ro), size=round(ro / 10), replace=False)
    for el in elements_to_mutate:
        pop[el] = swap_mutation(pop[el])

def particle_swarm_optimization(cities,
                                ro=30,
                                max_number_of_iterations=1000,
                                decrement_factor=0.975,
                                c1=2,  # cognitive param
                                c2=2,  # social param
                                inertia_weight=0.99,
                                wait_interval=100,
                                hybrid_evolutionary_approach=False):
    # For reproducibility
    random.seed(10)
    np.random.seed(10)

    not_primes_bool = generate_not_primes_numbers_array(cities)
    # iteration counter
    t = 0
    # Initialization
    path_length = len(cities) - 1 #Exclude 0 because it's a special node
    pop = create_pop_particles(pop_size=ro, particle_length=path_length)
    best_pop = pop
    vel = create_pop_velocities(pop_size=ro, particle_length=path_length)
    paths = generate_santas_path_from_particles_pop(pop)
    distances = evaluate_paths(paths, cities, not_primes_bool)
    best_distances = distances
    # Initialize support variables where best results will be stored
    best_iteration = t
    best_particle_index = np.argsort(distances)[0]
    best_particle = pop[best_particle_index]
    best_path = paths[best_particle_index]
    global_min_distance = distances[best_particle_index]
    r1 = np.random.uniform(low=0, high=1)
    r2 = np.random.uniform(low=0, high=1)

    results_for_each_iteration = [[t, global_min_distance]]
    for t in range(1, max_number_of_iterations + 1):
        if t % 100 == 0:
            print("Iteration: {}, minimum distance so far: {}".format(t, global_min_distance))
        # Update velocity and pop
        vel = vel + c1 * r1 * (best_pop - pop)
        vel = vel + c2 * r2 * (best_particle - pop)
        pop = pop + vel
        #Mutate 10% of pop elements
        mutate_pop_elements(pop, ro)
        # Compute new path and the related distances
        paths = generate_santas_path_from_particles_pop_complete_dataset(pop)
        distances = evaluate_paths(paths, cities, not_primes_bool)
        # Find which paths lead to an improved distance and update the corresponding values
        distance_comparison = distances < best_distances
        best_distances[distance_comparison] = distances[distance_comparison]
        best_pop[distance_comparison] = pop[distance_comparison]
        # Update global best results if necessary
        current_iteration_best_index = np.argsort(distances)[0]
        current_min_distance = distances[current_iteration_best_index]
        if current_min_distance < global_min_distance:
            best_iteration = t
            best_particle = pop[current_iteration_best_index]
            global_min_distance = current_min_distance
            best_path = paths[current_iteration_best_index]
        results_for_each_iteration.append([t, global_min_distance])

        #Interrupt the algorithm if the solution is no longer improving
        if t > best_iteration + wait_interval:
            break

        # Hybrid approach: update velocities and pop params of the worst elements
        if hybrid_evolutionary_approach:
            particle_index_ordered_by_distance = np.argsort(distances)
            best_particles_index = particle_index_ordered_by_distance[0:int(ro / 2)]
            worst_particles_index = particle_index_ordered_by_distance[int(ro / 2):]
            pop[worst_particles_index] = pop[best_particles_index]
            vel[worst_particles_index] = vel[best_particles_index]

        # Update inertia and velocity
        inertia_weight = inertia_weight * decrement_factor
        if inertia_weight < 0.4:
            inertia_weight = 0.4
        vel = inertia_weight * vel

    return results_for_each_iteration, best_path

def compute_final_paths(paths, clusters_els):
    """
    Compute the complete paths taking into consideration the cluster orders and how the dataset was clusterized
    :param paths: array containing the chosen cluster combinations
    :param clusters_els: array containing the cluster related to each dataset point
    :return:
    """
    end_paths = []
    clusters_els = np.array(clusters_els)
    for path in paths:
        current_path = []
        path = np.array(path)
        for cluster_path in clusters_els[path]:
            current_path = np.concatenate((current_path, cluster_path))
        current_path = current_path.astype(int)
        end_paths.append(current_path)
    return end_paths

def evaluate_paths_cluster(paths, cities):
    """
    Compute the euclidean distance for each path
    :param paths:
    :param cities:
    :return:
    """
    distances = []
    for path in paths:
        distances.append(santas_path.total_length_straight(path, cities))
    return np.array(distances)

def generate_santas_path_from_particles_cluster(population, cluster_elements):
    """
    Transform the given particles in feasible paths using the SVP rule
    :param population:
    :param cluster_elements:
    :return:
    """
    path_populations = []
    for particle in population:
        path = np.sort(cluster_elements)[np.argsort(particle)]
        # This city is the North Pole, starting and ending point of each path: it should not be considered
        path = np.delete(path, np.where(path == 0))
        path_populations.append(path)
    return path_populations

def initialize_clusters(cluster_dict, ro, cluster_indexes, cities):
    """
    Initialize the fundamental elements needed by PSO algorithm
    :param cluster_dict:
    :param ro:
    :param cluster_indexes:
    :param cities:
    :return:
    """
    clusters_pop = []
    clusters_vel = []
    clusters_best_path = []
    clusters_best_particle = []
    clusters_distances = []
    for c in cluster_dict.items():
        pop = create_pop_particles(pop_size=ro, particle_length=c[1])
        clusters_pop.append(pop)
        vel = create_pop_velocities(pop_size=ro, particle_length=c[1])
        clusters_vel.append(vel)
        paths = generate_santas_path_from_particles_cluster(pop, np.where(cluster_indexes == c[0])[0])
        distances = evaluate_paths_cluster(paths, cities)
        clusters_distances.append(distances)
        best_path_index = np.argsort(distances)[0]
        best_path = paths[best_path_index]
        clusters_best_path.append(best_path)

        clusters_best_particle.append(pop[best_path_index])
    return clusters_pop, clusters_vel, clusters_best_path, clusters_best_particle, clusters_distances

def initialize_overall_params(cities, clusters_best_path, clusters_pop, not_primes_bool, ro, number_of_clusters):
    """
    Initialize parameters needed by PSO for findingn the best order for assembling clusters
    :param cities:
    :param clusters_best_path:
    :param clusters_pop:
    :param not_primes_bool:
    :param ro:
    :return:
    """

    overall_pop = create_pop_particles(pop_size=ro, particle_length=number_of_clusters)
    overall_vel = create_pop_velocities(pop_size=ro, particle_length=number_of_clusters)
    overall_paths = generate_santas_path_from_particles_pop(overall_pop)
    final_complete_paths = compute_final_paths(overall_paths, clusters_best_path)
    final_complete_distances = evaluate_paths(paths=final_complete_paths, cities=cities,
                                              not_primes_bool=not_primes_bool)
    return final_complete_distances, final_complete_paths, overall_pop, overall_vel

def generate_cluster_dict(clusterized_cities):
    """
    Generate a dictionary where its keys are the clusters ID and its values are the number of its elements
    :param clusterized_cities:
    :return:
    """
    cluster_dict = {}
    for cluster_index in set(clusterized_cities):
        cluster_dict[cluster_index] = np.where(clusterized_cities == cluster_index)[0].shape[0]
    return cluster_dict

def generate_not_primes_numbers_array(cities):
    """
    Find which are the cities whose ID is a prime number
    :param cities:
    :return:
    """
    np_not_prime = np.vectorize(santas_path.not_prime)
    nums = np.arange(0, len(cities))
    not_primes_bool = np_not_prime(nums)
    return not_primes_bool

def cluster_particle_swarm_optimization(
        cities,
        clusterized_cities,
        ro=30,
        max_number_of_iterations=800,
        number_of_iterations_clusters=800,
        hybrid_evolutionary_approach=False,
        decrement_factor=0.975,
        c1=2,  # cognitive param
        c2=2,  # social param
        inertia_weight=0.99,
        wait_interval=100):
    # For reproducibility
    random.seed(10)
    np.random.seed(10)
    #Add list of prime numbers
    not_primes_bool = generate_not_primes_numbers_array(cities)
    t = 0

    original_inertia_weight = inertia_weight
    # Create elements necessary for performing PSO on each cluster
    cluster_dict = generate_cluster_dict(clusterized_cities)
    clusters_pop, clusters_vel, clusters_best_path, clusters_best_particle, clusters_distances = initialize_clusters(
        cluster_dict, ro, clusterized_cities, cities)
    clusters_best_pop = clusters_pop
    r1 = np.random.uniform(low=0, high=1)
    r2 = np.random.uniform(low=0, high=1)
    number_of_clusters = len(clusters_pop)

    for cluster in range(number_of_clusters):
        if cluster % 100 == 0:
            print("Start working on cluster {}".format(cluster))
        best_overall_iteration = 0
        inertia_weight = original_inertia_weight
        # Get initial values for cluster element
        current_pop = clusters_pop[cluster]
        current_best_pop = clusters_best_pop[cluster]
        current_vel = clusters_vel[cluster]
        current_dist = clusters_distances[cluster]
        current_best_particle = clusters_best_particle[cluster]
        current_best_path = clusters_best_path[cluster]
        for i in range(1, number_of_iterations_clusters + 1):

            # Perform mutation
            mutate_pop_elements(current_pop, ro)
            # Update velocity and pop
            current_vel = current_vel + c1 * r1 * (current_best_pop - current_pop)
            current_vel = current_vel + c2 * r2 * (current_best_particle - current_pop)
            current_pop = current_pop + current_vel
            # Transform the particle in feasible paths
            current_paths = generate_santas_path_from_particles_cluster(current_pop,
                                                                        np.where(clusterized_cities == cluster)[0])
            distances = evaluate_paths_cluster(current_paths, cities)

            # Update best particle if needed
            current_best_dist = np.sort(current_dist)[0]
            distance_comparison = distances < current_dist
            current_dist[distance_comparison] = distances[distance_comparison]
            current_best_pop[distance_comparison] = current_pop[distance_comparison]
            local_best_particle_distance = np.sort(distances)[0]
            if local_best_particle_distance < current_best_dist:
                best_particle_index = np.argsort(current_dist)[0]
                current_best_particle = current_pop[best_particle_index]
                current_best_path = current_paths[best_particle_index]
                best_overall_iteration = i

            # Stop doing PSO on the given cluster if the solution is no longer improving
            if i > best_overall_iteration + wait_interval:
                break

            if hybrid_evolutionary_approach:
                best_elements_indexes = np.argsort(distances)[0:int(ro / 2)]
                worst_elements_indexes = np.argsort(distances)[int(ro / 2):]
                current_pop[worst_elements_indexes] = current_pop[best_elements_indexes]
                current_vel[worst_elements_indexes] = current_vel[worst_elements_indexes]
            inertia_weight = inertia_weight * decrement_factor

            if inertia_weight < 0.4:
                inertia_weight = 0.4
            current_vel = current_vel * inertia_weight

        # Update storing structures
        clusters_best_pop[cluster] = current_best_pop
        clusters_distances[cluster] = current_dist
        clusters_pop[cluster] = current_pop
        clusters_vel[cluster] = current_vel
        clusters_best_particle[cluster] = current_best_particle
        clusters_best_path[cluster] = current_best_path
    end_pso_each_cluster = datetime.now()
    print("Ended PSO on each cluster")

    #Initialize elements useful for performing PSO on clusters order
    final_complete_distances, final_complete_paths, overall_pop, overall_vel = initialize_overall_params(
        cities, clusters_best_path, clusters_pop, not_primes_bool, ro, number_of_clusters)

    best_final_complete_distances = final_complete_distances
    best_overall_pop = overall_pop
    best_overall_particle = overall_pop[np.argsort(final_complete_distances)[0]]
    final_path = final_complete_paths[np.argsort(final_complete_distances)[0]]
    best_overall_iteration = t
    best_overall_distance = np.sort(final_complete_distances)[0]

    inertia_weight = original_inertia_weight
    results_for_each_iteration = [[t, best_overall_distance]]
    for t in range(1, max_number_of_iterations + 1):

        if t % 100 == 0:
            print("Iteration_overall: {}, min_distance: {}".format(t, best_overall_distance))
        # Update overall particles

        overall_vel = overall_vel + c1 * r1 * (best_overall_pop - overall_pop)
        overall_vel = overall_vel + c2 * r2 * (best_overall_particle - overall_pop)
        overall_pop = overall_pop + overall_vel
        overall_paths = generate_santas_path_from_particles_pop(overall_pop)
        final_complete_paths = compute_final_paths(overall_paths, clusters_best_path)
        final_complete_distances = evaluate_paths(paths=final_complete_paths, cities=cities,
                                                  not_primes_bool=not_primes_bool)
        # Find which paths lead to an improved distance and update the corresponding values
        distance_comparison = final_complete_distances < best_final_complete_distances
        best_final_complete_distances[distance_comparison] = final_complete_distances[distance_comparison]
        best_overall_pop[distance_comparison] = overall_pop[distance_comparison]
        # Update global best results if necessary
        current_iteration_best_index = np.argsort(final_complete_distances)[0]
        current_min_distance = final_complete_distances[current_iteration_best_index]
        if current_min_distance < best_overall_distance:
            best_overall_iteration = t
            best_overall_particle = overall_pop[current_iteration_best_index]
            final_path = final_complete_paths[current_iteration_best_index]
            best_overall_distance = current_min_distance

        results_for_each_iteration.append([t, best_overall_distance])
        if hybrid_evolutionary_approach:
            best_elements_indexes = np.argsort(final_complete_distances)[0:int(ro / 2)]
            worst_elements_indexes = np.argsort(final_complete_distances)[int(ro / 2):]
            overall_pop[worst_elements_indexes] = overall_pop[best_elements_indexes]
            overall_vel[worst_elements_indexes] = overall_vel[worst_elements_indexes]

        if t > best_overall_iteration + wait_interval:
            print("The algorithm is no longer improving")
            break

        inertia_weight = inertia_weight * decrement_factor
        if inertia_weight < 0.4:
            inertia_weight = 0.4
        overall_vel = inertia_weight * overall_vel

    return results_for_each_iteration, final_path, clusters_best_path, end_pso_each_cluster


# This file hold EXTRACT functions
from bs4 import BeautifulSoup

from fipe.scripts.loggers import get_logger

logger = get_logger(__name__)

def scrape_complete_tbody(driver, new_columns: list[str]) -> dict:
    """
    Extracts the HTML table within the tbody tags and returns the information as a dictionary.

    Args:
        driver (WebDriver): The WebDriver instance for controlling the browser.
        new_columns (list[str]): New columns coming from the YML configuration file.

    Returns:
        dict: A dictionary containing the extracted information. The keys are the columns from `new_columns`
              and the values are the corresponding values from the HTML table.
    """
    # Get the new URL
    soup = BeautifulSoup(driver.page_source, "html.parser")

    tbody = soup.find_all("tbody")

    row_tbody = tbody[0]
    td_infos = row_tbody.findChildren("td")

    keys_to_remove = [
        "Mês de referência:",
        "Código Fipe:",
        "Marca:",
        "Modelo:",
        "Ano Modelo:",
        "Autenticação",
        "Data da consulta",
        "Preço Médio",
    ]

    value_tab = [
        " ".join(tr.string.split())
        for tr in td_infos
        if tr.string not in keys_to_remove
    ]

    complete_info = dict(zip(new_columns, value_tab))

    return complete_info

def scrape_options_month_year(driver) -> list[str]:
    """
    Extracts all available Reference Months from the HTML.

    Args:
        driver (WebDriver): The WebDriver instance for controlling the browser.

    Returns:
        list[str]: A list of strings representing the available Reference Months.
    """

    logger.info("Extracting all Reference Months available")
    # I can pick the values AVAILABLE, or using Requests or BeautifulSoup
    soup = BeautifulSoup(driver.page_source, "html.parser")

    # Month and Year Available
    options_months_years = soup.find_all(
        "select", attrs={"data-tabref": "selectTabelaReferenciacarro"}
    )

    months_years_available = [
        " ".join(month_year.text.split())
        for month_year in options_months_years[0].contents
        if month_year.text != ""
    ]

    return months_years_available

def scrape_options_brands(driver) -> list[str]:
    """
    Extracts all available brands from the HTML.

    Args:
        driver (WebDriver): The WebDriver instance for controlling the browser.

    Returns:
        list[str]: A list of strings representing the available brands.
    """

    logger.info("Extracting all Brands available")
    # I can pick the values AVAILABLE, or using Requests or BeautifulSoup

    soup = BeautifulSoup(driver.page_source, "html.parser")

    # Brands Available
    options_brands = soup.find_all(
        "select",
        attrs={
            "data-placeholder": "Digite ou selecione a marca do veiculo",
            "data-tipo": "marca",
        },
    )

    brands = [
        brand.text.strip()
        for brand in options_brands[0].children
        if brand.text.strip() != "" and brand.text.strip() != " "
    ]

    return brands

def scrape_options_models(driver) -> list[str]:
    """
    Extract all Models available for the current Brand
    If Brand available is Nissan, then will get Sentra, Versa, Frontier and so on.
    Args:
        driver (WebDriver): The WebDriver instance for controlling the browser.

    Returns:
        list[str]: List of models.
    """

    logger.info("Extracting all models available")
    # Get the new URL
    soup = BeautifulSoup(driver.page_source, "html.parser")
    options_models_years = soup.find("div", attrs={"class": "step-2"})

    # Get ALL Models available according to a BRAND selected
    models_select = options_models_years.findChildren("select")[0]
    models = models_select.contents

    # My List of Models
    all_models = [model.text for model in models if model.text != ""]
    return all_models

def scrape_manufacturing_year_fuel(driver) -> list[str]:
    """
    Extracts the available MANUFACTURING YEAR - FUEL options.

    Args:
        driver (WebDriver): The WebDriver instance for controlling the browser.

    Returns:
        list[str]: A list of strings representing the available MANUFACTURING YEAR - FUEL options.
    """
    logger.info("Extracting all MANUFACTURING YEAR - FUEL available")
    # Get the new URL
    soup = BeautifulSoup(driver.page_source, "html.parser")

    year_fuel_for_especific_brand = soup.find(
        "select",
        attrs={"data-placeholder": "Digite ou selecione o ano modelo do veiculo"},
    )

    year_fuel_according_brand_and_model = [
        year.text for year in year_fuel_for_especific_brand if year.text != ""
    ]

    return year_fuel_according_brand_and_model


"""empty message

Revision ID: 2e7f50426ccf
Revises: a2f5a0132bde
Create Date: 2018-03-13 12:20:00.301645

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = '2e7f50426ccf'
down_revision = 'a2f5a0132bde'
branch_labels = None
depends_on = None

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('tag',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('name', sa.String(length=50), nullable=True),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_table('translation_tags',
    sa.Column('translation_id', sa.Integer(), nullable=False),
    sa.Column('tag_id', sa.Integer(), nullable=False),
    sa.ForeignKeyConstraint(['tag_id'], ['tag.id'], ),
    sa.ForeignKeyConstraint(['translation_id'], ['translation.id'], ),
    sa.PrimaryKeyConstraint('translation_id', 'tag_id')
    )
    # ### end Alembic commands ###

def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table('translation_tags')
    op.drop_table('tag')
    # ### end Alembic commands ###


import requests

def get_info():
    response = requests.get(url="https://yobit.net/api/3/info") # Получаем инфу по парам

    with open("output_data/Yobit_data/info.txt", "w") as file:
        file.write(response.text)

    return response.text

def get_ticker(coin1="btc", coin2="usd"): # Статистика за последние 24 часа
    response = requests.get(url=f"https://yobit.net/api/3/ticker/{coin1}_{coin2}?ignore_invalid=1") # /Перечисляем пару   ?ignore_invalid=1 если какой-то пары нет, то вместо ошибки выдаст то, что есть

    with open("output_data/Yobit_data/ticker.txt", "w") as file:
        file.write(response.text)

    return response.text

def get_depth(coin1="btc", coin2="usd", limit=50): # Список активных ордеров  limit - глубина для выводa(150 дефолт, 2000 макс)
    response = requests.get(url=f"https://yobit.net/api/3/depth/{coin1}_{coin2}?limit={limit}?ignore_invalid=1") # /Перечисляем пару   ?ignore_invalid=1 если какой-то пары нет, то вместо ошибки выдаст то, что есть

    with open("output_data/Yobit_data/depth.txt", "w") as file:
        file.write(response.text)

    """        Список предложений на покупку         """

    bids = response.json()[f"{coin1}_{coin2}"]["bids"] # Отсеиваем ненужное и оставляем нужный список предложений на покупку с ценой и количеством

    total_bids_amount = 0
    for item in bids: # (цена1 * количество1) + (цена2 * количество2) +... не понял для чего это
        price = item[0]
        coin_amount = item[1]
        total_bids_amount += price * coin_amount

    # Минимальная цена с количеством
    min_bid_price = bids[0][0]
    min_bid_amount = bids[0][1]

    '''        Список предложений на продажу         '''

    asks = response.json()[f"{coin1}_{coin2}"]["asks"] # Отсеиваем ненужное и оставляем нужный список предложений на продажу с ценой и количеством

    total_asks_amount = 0
    for item in asks: # (цена1 * количество1) + (цена2 * количество2) +... не понял для чего это
        price = item[0]
        coin_amount = item[1]
        total_asks_amount += price * coin_amount

    # Минимальная цена с количеством
    min_ask_price = asks[0][0]
    min_ask_amount = asks[0][1]

    # Визуальное оформление валюты результата
    if coin2 == "usd":
        currency = "$"
    elif coin2 == "rur":
        currency = "₽"

    return f"\nTotal bids: {total_bids_amount} {currency}\nMin bid price: {min_bid_price} {currency}\nAmount: {min_bid_amount}\n\nTotal asks: {total_asks_amount} {currency}\nMin ask price: {min_ask_price} {currency}\nAmount: {min_ask_amount}"

def get_trades(coin1="btc", coin2="usd", limit=150): # Инфа о последних сделках указанных пар
    response = requests.get(url=f"https://yobit.net/api/3/trades/{coin1}_{coin2}?limit={limit}?ignore_invalid=1") # /Перечисляем пару   ?ignore_invalid=1 если какой-то пары нет, то вместо ошибки выдаст то, что есть

    with open("output_data/Yobit_data/trades.txt", "w") as file:
        file.write(response.text)

    total_trade_ask = 0
    total_trade_bid = 0

    for item in response.json()[f"{coin1}_{coin2}"]:
        if item["type"] == "ask":
            total_trade_ask += item["price"] * item["amount"]
        else:
            total_trade_bid += item["price"] * item["amount"]

    info = f"[-] TOTAL {coin1} SELL: {round(total_trade_ask, 2)} $\n[+] TOTAL {coin1} BUY: {round(total_trade_bid, 2)} $"

    return info

def main():
    coin1 = input("Введите криптовалюту: ")
    coin2 = input("Введите желаемую валюту: ")
    #print(get_info())
    #print(get_ticker())
    #print(get_ticker(coin1="eth", coin2=coin2))
    print(get_depth(coin1=coin1, coin2=coin2))
    #print(get_depth(coin1="doge"))
    #print(get_depth(coin1="doge", limit=2000))
    #print(get_trades(coin1=coin1, coin2=coin2))

if __name__== "__main__":
    main()


import pandas as pd
import math
from scipy import spatial
from sklearn.feature_extraction.text import TfidfVectorizer
import pickle
import nltk
nltk.download('stopwords')
nltk.download('punkt')
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize

# ***************************************************************************************************************************************
# Query
#query = "aubergine fennel bulb red onion mushrooms olive oil garlic cloves"
#query = query.lower()
#q_list = query.split()
# print (q_list)

# ***************************************************************************************************************************************
# Read file
df = pd.read_csv("/Users/placid_brain/Documents/IR stuff/lsa/recipes_w_search_terms.csv",error_bad_lines=False, engine="python", nrows=1000)

# ***************************************************************************************************************************************
# Processing
def function(ini_list):
  new_cell = ini_list.strip('][').split(', ')
  for item in new_cell:
    item = item.replace("'","")
  return new_cell

df['ingredients'] = df['ingredients'].apply(function)
df['steps'] = df['steps'].apply(function)
df1 = df[['id', 'name','ingredients','steps']]
df1.to_pickle("./df1_926.pkl")
# print (df1)
# print(df1['ingredients'])
# print ((df1['ingredients'][0]))
# print ((df1['ingredients'][1]))

# ***************************************************************************************************************************************
# tf-idf vector for documents
data = []
for i in range (len(df1['ingredients'])):
    ingre_str = ''.join(df1['ingredients'][i])
    # print (ingre_str)
    data.append(ingre_str)

# print(data)

tfvec = TfidfVectorizer()
tdf = tfvec.fit_transform(data)
tf_idf_docs_matrix = pd.DataFrame(tdf.toarray(), columns = tfvec.get_feature_names())

'''term_list = tfvec.get_feature_names()
with open('list_pickle', 'wb') as fh:
  pickle.dump(term_list, fh)
tf_idf_docs_matrix.to_pickle("./tf_idf_docs_matrix.pkl") '''
# print(tf_idf_docs_matrix)

# get list of tf-idf vectors of recipes
# print (tf_idf_docs_matrix.values.tolist())

# ***************************************************************************************************************************************

def vector_gen(q_list):
  # tf-idf vector for query
  # print(tfvec.get_feature_names())

  query_tfidf_vector = []
  for i in tfvec.get_feature_names():
      # print(i)
      if i not in q_list:
          query_tfidf_vector.append(0.0)
      else:
          # append tf-idf of the term for query
          # print(i)
          query_tfidf_vector.append(tfvec.idf_[tfvec.vocabulary_[i]])
  return query_tfidf_vector

# print (query_tfidf_vector)
# print(len(query_tfidf_vector))

def search(query):
  sw = set(stopwords.words('english'))

  query = query.lower()

  query_tokens = word_tokenize(query)
  q_list = []

  for w in query_tokens:
      if w not in sw:
          q_list.append(w)
  query_tfidf_vector = vector_gen(q_list)

  # similarity - cosine similarity
  # dict to store cosine similarity between query and documents (key - doc index, value - score)
  similarity_list = []

  for i in range(len(tf_idf_docs_matrix.values.tolist())):
      # print(tf_idf_docs_matrix.values.tolist()[i])
      cosine_similarity = 1 - spatial.distance.cosine(query_tfidf_vector, tf_idf_docs_matrix.values.tolist()[i])
      similarity_list.append(cosine_similarity)

  # print(similarity_list)

  top_5_recipes = sorted(range(len(similarity_list)), key=lambda i: similarity_list[i], reverse=True)[:5]

  # print(top_5_recipes)

  for i in range(len(top_5_recipes)):
    # print(top_5_recipes[i])
    print("{}. {}".format(i+1, df1['name'][top_5_recipes[i]]))

#search('aubergine fennel bulb red onion mushrooms olive oil garlic cloves')

#pickle.dump(model, open("model", 'wb'))

class VSM:
  def __init__(self):
    #super(VSM, self).__init__()
    #self.query = query
    #self.search(self.query)
    #super(VSM, self).__init__()
    pass
  def vector_gen(self, q_list):
    # tf-idf vector for query
    # print(tfvec.get_feature_names())

    query_tfidf_vector = []
    for i in tfvec.get_feature_names():
        # print(i)
        if i not in q_list:
            query_tfidf_vector.append(0.0)
        else:
            # append tf-idf of the term for query
            # print(i)
            query_tfidf_vector.append(tfvec.idf_[tfvec.vocabulary_[i]])
    return query_tfidf_vector

  def search(self,query):
    query = query.lower()
    q_list = query.split()

    if set(q_list).issubset(set(tfvec.get_feature_names()))== True:
      query_tfidf_vector = self.vector_gen(q_list)

      # similarity - cosine similarity
      # dict to store cosine similarity between query and documents (key - doc index, value - score)
      similarity_list = []
      final_list = []

      for i in range(len(tf_idf_docs_matrix.values.tolist())):
          # print(tf_idf_docs_matrix.values.tolist()[i])
          cosine_similarity = 1 - spatial.distance.cosine(query_tfidf_vector, tf_idf_docs_matrix.values.tolist()[i])
          similarity_list.append(cosine_similarity)

      # print(similarity_list)

      top_5_recipes = sorted(range(len(similarity_list)), key=lambda i: similarity_list[i], reverse=True)[:5]

      for i in range(len(top_5_recipes)):
        final_list.append(df1['name'][top_5_recipes[i]])
        print("{}. {}".format(top_5_recipes[i], df1['name'][top_5_recipes[i]]))
      return final_list
    else:
      ["No matching results found"]

vsm = VSM()
vsm.search('aubergine fennel bulb red onion mushrooms olive oil garlic cloves')

with open('model_pickle_926','wb') as f:
  pickle.dump(vsm,f)


# Author: Matthew Wicker

from statsmodels.stats.proportion import proportion_confint
import math
import numpy as np
import tensorflow as tf
from tqdm import trange
from . import attacks

def propagate_interval(W, b, x_l, x_u, marg=0, b_marg=0):
    marg = tf.divide(marg, 2)
    b_marg = tf.divide(b_marg, 2)

    x_mu = tf.divide(tf.math.add(x_u, x_l), 2)
    x_r =  tf.divide(tf.math.subtract(x_u, x_l), 2)

    W_mu = tf.cast(W, dtype=tf.float64)
    W_r =  tf.cast(marg, dtype=tf.float64)

    b = tf.cast(b, dtype=tf.float64)
    b_marg = tf.cast(b_marg, dtype=tf.float64)
    b_u =  tf.cast(b + b_marg, dtype=tf.float64)
    b_l =  tf.cast(b - b_marg, dtype=tf.float64)

    h_mu = tf.matmul(x_mu, W_mu)
    x_rad = tf.matmul(x_r, tf.math.abs(W_mu))
    try:
        W_rad = tf.matmul(tf.abs(x_mu), W_r)
        Quad = tf.matmul(tf.abs(x_r), tf.abs(W_r))
    except:
        W_rad = 0 # This happens if margin stays zero
        Quad = 0

    h_u = tf.add(tf.add(tf.add(tf.add(h_mu, x_rad), W_rad), Quad), b_u)
    h_l = tf.add(tf.subtract(tf.subtract(tf.subtract(h_mu, x_rad), W_rad), Quad), b_l)

    return h_l, h_u

def propagate_interval_exact(W, b, x_l, x_u, marg=0, b_marg=0):
    """
    Function which does matrix multiplication but with weight and
    input intervals.
    """
    x_l = tf.cast(x_l, dtype=tf.float32);x_u = tf.cast(x_u, dtype=tf.float32)
    W = tf.cast(W, dtype=tf.float32); b = tf.cast(b, dtype=tf.float32)
    marg = tf.cast(marg, dtype=tf.float32); b_marg = tf.cast(b_marg, dtype=tf.float32)
    x_l = tf.squeeze(x_l); x_u = tf.squeeze(x_u)
    W_l, W_u = W-marg, W+marg           #Use eps as small symetric difference about the mean
    b_l, b_u = b-b_marg, b+b_marg       #Use eps as small symetric difference about the mean
    h_max = np.zeros(len(W[0]))         #Placeholder variable for return value
    h_min = np.zeros(len(W[0]))         #Placeholder variable for return value
    for i in range(len(W)):             #This is literally just a step-by-step matrix multiplication
        for j in range(len(W[0])):      # where we are taking the min and max of the possibilities
            out_arr = [W_l[i][j]*x_l[i], W_l[i][j]*x_u[i],
                       W_u[i][j]*x_l[i], W_u[i][j]*x_u[i]]
            h_min[j] += min(out_arr)
            h_max[j] += max(out_arr)
    h_min = h_min + b_l
    h_max = h_max + b_u
    return h_min, h_max         #Return the min and max of the intervals.
                                #(dont forget to apply activation function after)

def IBP_state(model, s0, s1, weights, weight_margin=0, logits=True, exact=False):
    h_l = s0
    h_u = s1
    layers = model.model.layers
    offset = 0
    for i in range(len(layers)):
        if(len(layers[i].get_weights()) == 0):
            h_u = model.model.layers[i](h_u)
            h_l = model.model.layers[i](h_l)
            offset += 1
            continue
        w, b = weights[2*(i-offset)], weights[(2*(i-offset))+1]
        sigma = model.posterior_var[2*(i-offset)]
        b_sigma = model.posterior_var[2*(i-offset)+1]
        marg = weight_margin*sigma
        b_marg = weight_margin*b_sigma
        if(exact):
            h_l, h_u = propagate_interval_exact(w, b, h_l, h_u, marg=marg, b_marg=b_marg)
        else:
            h_l, h_u = propagate_interval(w, b, h_l, h_u, marg=marg, b_marg=b_marg)
        h_l = model.model.layers[i].activation(h_l)
        h_u = model.model.layers[i].activation(h_u)
    return h_l, h_u

def IBP_full_multiproc(args):
    actives, s0, s1, weights, weight_margin, predicate, posterior_var = args
    h_l = s0
    h_u = s1
    layers = int(len(posterior_var)/2)
    offset = 0
    for i in range(layers):
        w, b = weights[2*(i-offset)], weights[(2*(i-offset))+1]
        sigma = posterior_var[2*(i-offset)]
        b_sigma = posterior_var[2*(i-offset)+1]
        marg = weight_margin*sigma
        b_marg = weight_margin*b_sigma
        h_l, h_u = propagate_interval(w, b, h_l, h_u, marg=marg, b_marg=b_marg)
        #print("Layer %s Bounds: "%(i), h_l, h_u)
        h_l = actives[i](h_l) #model.model.layers[i].activation(h_l)
        h_u = actives[i](h_u) #model.model.layers[i].activation(h_u)

    ol, ou = h_l, h_u
    lower = np.squeeze(s0)[0:len(ol)] + ol;
    upper = np.squeeze(s1)[0:len(ou)] + ou

    if(predicate(np.squeeze(s0), np.squeeze(s1), np.squeeze(lower), np.squeeze(upper))):
        ol = np.squeeze(ol); ou = np.squeeze(ou)
        return [lower, upper]
    else:
        return None

def IBP_conf(model, s0, s1, weights, weight_margin=0, logits=True):
    h_l = s0
    h_u = s1
    layers = model.model.layers
    offset = 0
    for i in range(len(layers)):
        if(len(layers[i].get_weights()) == 0):
            h_u = model.model.layers[i](h_u)
            h_l = model.model.layers[i](h_l)
            offset += 1
            continue
        if(i == len(layers)-1):
            # iterate over the number of classes:
            softmax_diffs = []
            num_classes = np.asarray(weights[(2*(i-offset))+1]).shape[-1]
            print("number of classes: ", num_classes)
            for k in range(num_classes):
                class_diffs = []
                for l in range(num_classes):
                    diff = weights[2*(i-offset)][:,k] - weights[2*(i-offset)][:,l]
                    max_diff = np.maximum(diff, 0)
                    min_diff = np.minimum(diff, 0)
                    bias_diff = weights[(2*(i-offset))+1][k] - weights[(2*(i-offset))+1][l]
                    #logit_diff = (max_diff * h_u) + (min_diff * h_l)  + bias_diff
                    logit_diff = np.sum(max_diff * h_u) + np.sum(min_diff * h_l)  + bias_diff
                    class_diff = logit_diff
                    #print(class_diff)
                    #if(class_diff == 0.0):
                    #    print("These should be the same ", k,l)
                    class_diffs.append(class_diff)
                diff_val = -np.log(np.sum(np.exp(-1*np.asarray(class_diffs))))
                #diff_val = np.max(class_diffs)
                #diff_val = np.max(diff_val)
                print("Diff value for class ", k, " is ", diff_val)
                softmax_diffs.append(diff_val)
            return softmax_diffs
        w, b = weights[2*(i-offset)], weights[(2*(i-offset))+1]
        sigma = model.posterior_var[2*(i-offset)]
        marg = weight_margin*sigma
        h_l, h_u = propagate_interval(w, b, h_l, h_u, marg=marg)
    return None #error state if we hit this!

def IBP_learning(model, inp, weights, eps, predict=False):
    h_u = tf.clip_by_value(tf.math.add(inp, eps), 0.0, 1.0)
    h_l = tf.clip_by_value(tf.math.subtract(inp, eps), 0.0, 1.0)
    layers = model.model.layers
    offset = 0
    for i in range(len(layers)):
        if(len(layers[i].get_weights()) == 0):
            h_u = model.model.layers[i](h_u)
            h_l = model.model.layers[i](h_l)
            offset += 1
            continue
        w, b = weights[2*(i-offset)], weights[(2*(i-offset))+1]
        h_l, h_u = propagate_interval(w, b, h_l, h_u)
        if(i < len(layers)-1):
            h_l = model.model.layers[i].activation(h_l)
            h_u = model.model.layers[i].activation(h_u)
    return h_l, h_u

def pIBP(model, inp_l, inp_u, weights, predict=False):
    #if(predict == False):
    #    h_u = tf.clip_by_value(tf.math.add(inp, eps), 0.0, 1.0)
    #    h_l = tf.clip_by_value(tf.math.subtract(inp, eps), 0.0, 1.0)
    #else:
    h_u = inp_u
    h_l = inp_l
    layers = model.model.layers
    offset = 0
    for i in range(len(layers)):
        if(len(layers[i].get_weights()) == 0):
            h_u = model.model.layers[i](h_u)
            h_l = model.model.layers[i](h_l)
            offset += 1
            continue
        w, b = weights[2*(i-offset)], weights[(2*(i-offset))+1]
        if(len(w.shape) == 2):
            h_l, h_u = propagate_interval(w, b, h_l, h_u)
            activate = True
        elif(len(w.shape) == 4):
            h_l, h_u = propagate_conv2d(w, b, h_l, h_u)
            activate = True
        if(predict == False and i >= len(layers)-1):
            continue
        else:
            h_l = model.model.layers[i].activation(h_l)
            h_u = model.model.layers[i].activation(h_u)
    return h_l, h_u

# Code for merging overlapping intervals. Taken from here:
# https://stackoverflow.com/questions/49071081/merging-overlapping-intervals-in-python
# This function simple takes in a list of intervals and merges them into all
# continuous intervals and returns that list
def merge_intervals(intervals):
    sorted_intervals = sorted(intervals)
    interval_index = 0
    intervals = np.asarray(intervals)
    for  i in sorted_intervals:
        if i[0] > sorted_intervals[interval_index][1]:
            interval_index += 1
            sorted_intervals[interval_index] = i
        else:
            sorted_intervals[interval_index] = [sorted_intervals[interval_index][0], i[1]]
    return sorted_intervals[:interval_index+1]

"""
Given a set of disjoint intervals, compute the probability of a random
sample from a guassian falling in these intervals. (Taken from lemma)
of the document
"""
import math
from scipy.special import erf
def compute_erf_prob(intervals, mean, var):
    prob = 0.0
    for interval in intervals:
        #val1 = erf((mean-interval[0])/(math.sqrt(2)*(var)))
        #val2 = erf((mean-interval[1])/(math.sqrt(2)*(var)))
        val1 = erf((mean-interval[0])/(math.sqrt(2*var)))
        val2 = erf((mean-interval[1])/(math.sqrt(2*var)))
        prob += 0.5*(val1-val2)
    return prob

"""
Given a set of possibly overlapping intervals:
    - Merge all intervals into maximum continuous, disjoint intervals
    - compute probability of these disjoint intervals
    - do this for ALL values in a weight matrix
"""
def compute_interval_probs_weight(vector_intervals, marg, mean, var):
    means = mean; # vars = var
    prob_vec = np.zeros(vector_intervals[0].shape)
    for i in trange(len(vector_intervals[0])):
        for j in range(len(vector_intervals[0][0])):
            intervals = []
            for num_found in range(len(vector_intervals)):
                interval = [vector_intervals[num_found][i][j]-(var[i][j]*marg), vector_intervals[num_found][i][j]+(var[i][j]*marg)]
                intervals.append(interval)
            p = compute_erf_prob(merge_intervals(intervals), means[i][j], var[i][j])
            prob_vec[i][j] = p
    return np.asarray(prob_vec)

def compute_interval_probs_weight_m(arg):
    vector_intervals, marg, mean, var = arg
    means = mean; # vars = var
    prob_vec = np.zeros(vector_intervals[0].shape)
    for i in trange(len(vector_intervals[0])):
        for j in range(len(vector_intervals[0][0])):
            intervals = []
            for num_found in range(len(vector_intervals)):
                interval = [vector_intervals[num_found][i][j]-(var[i][j]*marg), vector_intervals[num_found][i][j]+(var[i][j]*marg)]
                intervals.append(interval)
            p = compute_erf_prob(merge_intervals(intervals), means[i][j], var[i][j])
            prob_vec[i][j] = p
    return np.asarray(prob_vec)

"""
Given a set of possibly overlapping intervals:
    - Merge all intervals into continuous, disjoint intervals
    - compute probability of these disjoint intervals
    - do this for ALL values in a *flat* bias matrix (vector)
"""
def compute_interval_probs_bias(vector_intervals, marg, mean, var):
    means = mean; #stds = var
    prob_vec = np.zeros(vector_intervals[0].shape)
    for i in range(len(vector_intervals[0])):
        intervals = []
        for num_found in range(len(vector_intervals)):
            #!*! Need to correct and make sure you scale margin
            interval = [vector_intervals[num_found][i]-(var[i]*marg), vector_intervals[num_found][i]+(var[i]*marg)]
            intervals.append(interval)
        p = compute_erf_prob(merge_intervals(intervals), means[i], var[i])
        prob_vec[i] = p
    return np.asarray(prob_vec)

def compute_interval_probs_bias_m(arg):
    vector_intervals, marg, mean, var = arg
    means = mean; #stds = var
    prob_vec = np.zeros(vector_intervals[0].shape)
    for i in range(len(vector_intervals[0])):
        intervals = []
        for num_found in range(len(vector_intervals)):
            #!*! Need to correct and make sure you scale margin
            interval = [vector_intervals[num_found][i]-(var[i]*marg), vector_intervals[num_found][i]+(var[i]*marg)]
            intervals.append(interval)
        p = compute_erf_prob(merge_intervals(intervals), means[i], var[i])
        prob_vec[i] = p
    return np.asarray(prob_vec).tolist()

def compute_probability(model, weight_intervals, margin, verbose=True, n_proc=10):
    full_p = 1.0
    if(verbose == True):
        func = range
    else:
        func = range
    args_bias = []
    args_weights = []
    for i in func(len(model.posterior_mean)):
        if(i % 2 == 0): # then its a weight vector
            #p = compute_interval_probs_weight(weight_intervals[i], margin, model.posterior_mean[i], np.asarray(model.posterior_var[i]))
            args_weights.append((weight_intervals[i], margin, model.posterior_mean[i], np.asarray(model.posterior_var[i])))
        else: # else it is a bias vector
            #p = compute_interval_probs_bias(weight_intervals[i], margin, model.posterior_mean[i], np.asarray(model.posterior_var[i]))
            args_bias.append((weight_intervals[i], margin, model.posterior_mean[i], np.asarray(model.posterior_var[i])))

    from multiprocessing import Pool
    proc_pool = Pool(n_proc)
    ps_bias = proc_pool.map(compute_interval_probs_bias_m, args_bias)
    ps_weight = proc_pool.map(compute_interval_probs_weight_m, args_weights)
    proc_pool.close()
    proc_pool.join()

    import itertools
    ps_bias = np.concatenate(ps_bias).ravel()
    ps_weight = np.asarray(list(itertools.chain(*(itertools.chain(*ps_weight)))))

    full_p *= np.prod(ps_bias)
    full_p *= np.prod(ps_weight)
    return full_p

def IBP_prob(model, s0, s1, w_marg, samples, predicate, i0=0, inflate=1.0):
    #w_marg = w_marg**2
    safe_weights = []
    safe_outputs = []
    for i in range(samples):
        model.model.set_weights(model.sample(inflate=inflate))
        ol, ou = IBP_state(model, s0, s1, model.model.get_weights(), w_marg)
        if(predicate(np.squeeze(s0), np.squeeze(s1), np.squeeze(ol), np.squeeze(ou))):
            safe_weights.append(model.model.get_weights())
            ol = np.squeeze(ol); ou = np.squeeze(ou)
            #lower = np.squeeze(s0)[0:len(ol)] + ol; upper = np.squeeze(s1)[0:len(ou)] + ou
            safe_outputs.append([-1,1]) # This is used ONLY for control loops which needs its own verification section
    print("Found %s safe intervals"%(len(safe_weights)))
    if(len(safe_weights) < 2):
        return 0.0, -1
    p = compute_probability(model, np.swapaxes(np.asarray(safe_weights),1,0), w_marg**2)
    return p, np.squeeze(safe_outputs)

def propagate(model, s0, s1, w_marg=0.0):
    #model.model.set_weights(w)
    ol, ou = IBP_state(model, s0, s1, model.model.get_weights(), w_marg)
    lower = np.squeeze(s0)[0:len(ol)] + ol;
    upper = np.squeeze(s1)[0:len(ou)] + ou
    return lower, upper

def IBP_fixed_w(model, s0, s1, w_marg, w, predicate, i0=0):
    model.model.set_weights(w)
    ol, ou = IBP_state(model, s0, s1, w, w_marg)
    lower = np.squeeze(s0)[0:len(ol)] + ol;
    upper = np.squeeze(s1)[0:len(ou)] + ou
    if(predicate(np.squeeze(s0), np.squeeze(s1), np.squeeze(lower), np.squeeze(upper))):
        #p = compute_probability(model, np.swapaxes(np.asarray([w]),1,0), w_marg)
        return 1.0, -1
    else:
        return 0.0, -1

def IBP_prob_dyn(model, s0, s1, w_marg, samples, predicate, i0=0, inflate=1.0, full_state=False):
    w_marg = w_marg**2
    safe_weights = []
    safe_outputs = []

    for i in trange(samples):
        model.model.set_weights(model.sample(inflate=inflate))
        ol, ou = IBP_state(model, s0, s1, model.model.get_weights(), w_marg)

        lower = np.squeeze(s0)[0:len(ol)] + ol;
        upper = np.squeeze(s1)[0:len(ou)] + ou

        if(predicate(np.squeeze(s0), np.squeeze(s1), np.squeeze(lower), np.squeeze(upper))):
            safe_weights.append(model.model.get_weights())
            ol = np.squeeze(ol); ou = np.squeeze(ou)
            safe_outputs.append([lower, upper]) # This is used ONLY for control loops which needs its own verification section

    print("Found %s safe intervals"%(len(safe_weights)))
    if(len(safe_weights) < 2):
        return 0.0, -1
    p = compute_probability(model, np.swapaxes(np.asarray(safe_weights),1,0), w_marg)
    return p, np.squeeze(safe_outputs)

def IBP_prob_dyn_m(model, s0, s1, w_marg, samples, predicate, i0=0, inflate=1.0, full_state=False, n_proc=20):
    w_marg = w_marg**2
    safe_weights = []
    safe_outputs = []

    actives = []
    layers = model.model.layers
    offset = 0
    for i in range(len(layers)):
        actives.append(model.model.layers[i].activation)

    #actives, s0, s1, weights, weight_margin
    args = []
    weights = []
    for i in range(samples):
        #model.model.set_weights(model.sample(inflate=inflate))
        w = model.sample(inflate=inflate)
        #actives, s0, s1, weights, weight_margin, predicate
        arg = (actives, s0, s1, w, w_marg, predicate, model.posterior_var)
        weights.append(w)
        args.append(arg)

    weights = np.asarray(weights)

    from multiprocessing import Pool
    proc_pool = Pool(n_proc)
    #IBP_full_multiproc
    safe_outputs = proc_pool.map(IBP_full_multiproc, args)
    proc_pool.close()
    proc_pool.join()

    none_indexes = []; ind = 0
    safe_outs = []
    for i in safe_outputs:
        if(i is not None):
            none_indexes.append(ind)
            safe_outs.append(i)
        ind += 1
    safe_weights = weights[none_indexes]
    safe_outputs = safe_outs #safe_outputs[none_indexes]

    print("Found %s safe intervals"%(len(safe_weights)))
    if(len(safe_weights) < 2):
        return 0.0, -1
    p = compute_probability(model, np.swapaxes(np.asarray(safe_weights),1,0), w_marg)
    return p, np.squeeze(safe_outputs)

def IBP_prob_w(model, s0, s1, w_marg, w, predicate, i0=0):
    model.model.set_weights(model.sample())
    ol, ou = IBP_state(model, s0, s1, w, w_marg)
    if(predicate(np.squeeze(s0)[i0:i0+2], np.squeeze(s1)[i0:i0+2], np.squeeze(ol)[i0:i0+2], np.squeeze(ou)[i0:i0+2])):
        p = compute_probability(model, np.swapaxes(np.asarray([w]),1,0), w_marg)
        return p, -1
    else:
        return 0.0, -1

# also known as the chernoff bound
def okamoto_bound(epsilon, delta):
    return (-1*.5) * math.log(float(delta)/2) * (1.0/(epsilon**2))

# This is h_a in the paper
def absolute_massart_halting(succ, trials, I, epsilon, delta, alpha):
    gamma = float(succ)/trials
    if(I[0] < 0.5 and I[1] > 0.5):
        return -1
    elif(I[1] < 0.5):
        val = I[1]
        h = (9/2.0)*(((3*val + epsilon)*(3*(1-val)-epsilon))**(-1))
        return math.ceil((h*(epsilon**2))**(-1) * math.log((delta - alpha)**(-1)))
    elif(I[0] >= 0.5):
        val = I[0]
        h = (9/2.0)*(((3*(1-val) + epsilon)*((3*val)+epsilon))**(-1))
        return math.ceil((h*(epsilon**2))**(-1) * math.log((delta - alpha)**(-1)))

"""

"""
def chernoff_bound_verification(model, inp, eps, cls, **kwargs):
    from tqdm import trange
    delta = kwargs.get('delta', 0.3)
    alpha = kwargs.get('alpha', 0.05)
    confidence = kwargs.get('confidence', 0.95)
    verbose = kwargs.get('verbose', False)
    epsilon = 1-confidence
    chernoff_bound = math.ceil( (1/(2*epsilon**2)) * math.log(2/delta) )
    softmax = 0
    for i in trange(chernoff_bound, desc="Sampling for Chernoff Bound Satisfaction"):
        model.set_weights(model.sample())
        logit_l, logit_u = IBP(model, inp, model.model.get_weights(), eps, predict=False)
        v1 = tf.one_hot(cls, depth=10)
        v2 = 1 - tf.one_hot(cls, depth=10)
        v1 = tf.squeeze(v1); v2 = tf.squeeze(v2)
        worst_case = tf.math.add(tf.math.multiply(v2, logit_u), tf.math.multiply(v1, logit_l))
        if(type(softmax) == int):
            softmax = model.model.layers[-1].activation(worst_case)
        else:
            softmax += model.model.layers[-1].activation(worst_case)
    return softmax
    #print("Not yet implimented")

"""
property - a function that takes a vector and returns a boolean if it was successful
"""
def massart_bound_check(model, inp, eps, cls, **kwargs):
    delta = kwargs.get('delta', 0.3)
    alpha = kwargs.get('alpha', 0.05)
    confidence = kwargs.get('confidence', 0.95)
    verbose = kwargs.get('verbose', False)

    atk_locs = []
    epsilon = 1-confidence
    chernoff_bound = math.ceil( (1/(2*epsilon**2)) * math.log(2/delta) )
    print("BayesKeras. Maximum sample bound = %s"%(chernoff_bound))
    successes, iterations, misses = 0.0, 0.0, 0.0
    halting_bound = chernoff_bound
    I = [0,1]
    while(iterations <= halting_bound):
        if(iterations > 0 and verbose):
            print("Working on iteration: %s \t Bound: %s \t Param: %s"%(iterations, halting_bound, successes/iterations))
        model.set_weights(model.sample())
        logit_l, logit_u = IBP(model, inp, model.model.get_weights(), eps, predict=False)
        v1 = tf.one_hot(cls, depth=10)
        v2 = 1 - tf.one_hot(cls, depth=10)
        worst_case = tf.math.add(tf.math.multiply(v2, logit_u), tf.math.multiply(v1, logit_l))
        if(np.argmax(np.squeeze(worst_case)) != cls):
            misses += 1
            result = 0
        else:
            result = 1
        successes += result
        iterations += 1
        # Final bounds computation below
        lb, ub = proportion_confint(successes, iterations, method='beta')
        if(math.isnan(lb)):
            lb = 0.0 # Setting lb to zero if it is Nans
        if(math.isnan(ub)):
            ub = 1.0 # Setting ub to one if it is Nans
        I = [lb, ub]
        hb = absolute_massart_halting(successes, iterations, I, epsilon, delta, alpha)
        if(hb == -1):
            halting_bound = chernoff_bound
        else:
            halting_bound = min(hb, chernoff_bound)
    if(verbose):
        print("Exited becuase %s >= %s"%(iterations, halting_bound))
    return successes/iterations
    #return None

def chernoff_model_robustness(model, input_0, input_1=None, verify=True, **kwargs):
    """
    Function to compute the model robustness of a given Bayesian posterior
    with statistically guarenteed precision via Chernoff concentration
    """
    return None

def chernoff_decision_robustness(model, input_0, input_1=None, verify=True, **kwargs):
    """
    Function to compute the decision robustness of a given Bayesian posterior
    with statistically guarenteed precision via Chernoff concentration
    """
    return None

def massart_model_robustness(model, input_0, input_1=None, verify=True, **kwargs):
    """
    Function to compute the model robustness of a given Bayesian posterior
    with statistically guarenteed precision via Chernoff concentration
    """
    return None

def massart_decision_robustness(model, input_0, input_1=None, verify=True, **kwargs):
    """
    Function to compute the decision robustness of a given Bayesian posterior
    with statistically guarenteed precision via Chernoff concentration
    """
    return None

def model_robustness_lower(model, input_0, input_1=None, verify=True, **kwargs):
    """
    Function to compute a guarenteed probabalistic lower bound on the model robustness
    of a Bayesian posterior distribution
    """
    return None

def model_robustness_upper(model, input_0, input_1=None, verify=True, **kwargs):
    """
    Function to compute a guarenteed probabalistic lower bound on the model robustness
    of a Bayesian posterior distribution
    """
    return None

def decision_robustness_lower(model, input_0, input_1=None, verify=True, **kwargs):
    """
    Function to compute a guarenteed probabalistic lower bound on the model robustness
    of a Bayesian posterior distribution
    """
    return None

def decision_robustness_upper(model, input_0, input_1=None, verify=True, **kwargs):
    """
    Function to compute a guarenteed probabalistic lower bound on the model robustness
    of a Bayesian posterior distribution
    """
    return None

def log_confidence_upper(model, input_0, input_1=None, verify=True, **kwargs):
    """
    Function to compute an upper bound on the log confidence of a model decision from a
    BNN
    """
    return None