# -*- coding: utf-8 -*-
"""
Created on Wed Jul 22 15:11:07 2020

@author: hca
"""
import numpy as np
import numpy.ma as ma
import pandas as pd
import os
import pyemma
import matplotlib.pyplot as plt
from matplotlib import colors
import matplotlib.pylab as pl
from matplotlib.ticker import FormatStrFormatter, AutoLocator, MaxNLocator, ScalarFormatter, LogLocator, AutoMinorLocator
import matplotlib.gridspec as gridspec

import Discretize

def plot_scheme(states_scheme):

    color_rect = ['white', 'dimgrey']
    for idx_combination, (label,combination) in enumerate(states_scheme.items()):
        fig, ax_scheme = plt.subplots(1,1, figsize=(5,1))
        print(label)

        #ax_scheme.s
        for idx,color_idx in enumerate(combination):
            color = color_rect[color_idx]
            if idx == 0 and color_idx == 1:
                color = 'darkorange'
            rect = plt.Rectangle((idx,0), 1, 1, facecolor=color, edgecolor='black', linewidth=5)
            #print(idx, color, rect)
            ax_scheme.add_patch(rect)
        #ax_scheme.set_title(label)
        ax_scheme.set_xlim([0,5])
        ax_scheme.axis('off')
        #ax_scheme.set_ylim([0,9])
        fig.show()
        fig.savefig(f'/media/dataHog/hca/msAcT-acylOctamer/results/figures/comb_states/{label}.png', dpi=300)
            #(idx, idx+1), 2, 1, linewidth=1, edgecolor='r', facecolor='none')
    #ax_scheme.set_xticks(ax_scheme.get_xticks()-0.5)
    #ax_scheme.set_xticklabels(['A', 'T', 'E', 'S', 'B'])

    #ax_scheme.set_ylim([0,9])
    #ax_scheme.axis('scaled')

class Plots:

    loc1, loc2 =0, 350
    plt.rc('font', size=10)
    p_type = {'normal' : {'grid_layout' : (3,1), 'bar_width' : 0.8, 'outer_grid_heights' : [8,4,8], 'figsize' : (11,11)},
          'inhibition' : {'grid_layout' : (2,1), 'bar_width' : 0.4,  'outer_grid_heights' : [6,4], 'figsize' :(10,11)},
          'water' : {'grid_layout' : (3,3), 'bar_width' : 0.4,  'outer_grid_heights' : [6,2,6], 'figsize' :(10,11)}}

    def __init__(self,
                 supra_project : dict,
                    dG : pd.DataFrame,
                    combinatorial : pd.DataFrame,
                    plot_specs :dict,
                    input_states : dict,
                    subset_states : dict,
                    labels_regions = ['A', 'T', 'E', 'S', 'B'],
                    mol_water={} ,
                    color2='darkorange',
                    mol2={'BuOH' : 'ViAc', 'BeOH' : 'BeAc'},
                    discretized_matrixes={},
                    discretized_trajectories={},
                    metrics= False,
                    metric_method='Mutual Information (normalized)'):

        self.supra_project : dict = supra_project
        self.dG : pd.DataFrame = dG
        self.combinatorial : pd.DataFrame = combinatorial
        self.plot_specs : dict = plot_specs
        self.input_states : dict = input_states
        self.subset_states : dict = subset_states
        self.mol2 :dict = mol2
        self.color2 : str = color2
        self.mol_water : dict = mol_water
        self.metrics :bool = metrics
        self.metric_method : str = metric_method
        if self.metrics:
            self.discretized_trajectories : dict = discretized_trajectories
            self.discretized_matrixes : dict = discretized_matrixes

    @classmethod
    def MsAcTPublicationFigure(cls):
        print(cls)
        return cls.plotBindingProfileAndCombinatorial(cls)

        #return cls.plotBindingProfileAndCombinatorial()

    def plot_dg(self, input_df):

        it = self.it

        dg, dg_m, dg_p = self.getDeltaG(input_df, it, self.idx_it)

        ax = self.ax_dG
        loc1, loc2 = self.loc1, self.loc2
        ranges = dg.iloc[loc1:loc2].index.values
        if it.split('_')[-1] == '5mM':
            it = f'{it.split("_")[0]} + {" ".join(it.split("_")[1:])}'

        if self.project_type == 'water' and self.mol != 'H2O':
            #loc2 = loc2 - 60
            ranges = dg.iloc[loc1:loc2].index.values
            dg_diff = dg.iloc[loc1:loc2].values - self.dG_base_df.iloc[loc1:loc2].values

            ln = ax.plot(ranges,
                         dg_diff,
                         linestyle=self.linestyle,
                         color= self.color_mol[self.idx_it],
                         label=it)
            if self.idx_mol == 0:
                ax.set_ylabel(r'$\Delta$$\Delta$G ($\it{k}$$_B$T)')
            else:
                ax.set_yticklabels([])
            ax.set_ylim(-3, 1)
            ax.legend(loc='lower right')
        else:
            ln = ax.plot(ranges,
                         dg.iloc[loc1:loc2],
                         '.',
                         color=self.color_mol[self.idx_it],
                         label=it) #, marker=marker, fillstyle=fillstyle)
            ax.fill_between(ranges, dg_m.iloc[loc1:loc2], dg_p.iloc[loc1:loc2], alpha=0.5, color=self.color_mol[self.idx_it])
            ax.set_ylabel(r'$\Delta$G ($\it{k}$$_B$T)')
            ax.legend(loc='upper right')
        #ax.set_ylim(-1,14)
        ax.set_xticks([])

        ax.set_xscale('log')
        ax.axhline(y=0, ls='--', color='darkgray')

        ax.set_title(self.mol)
        ax.yaxis.set_minor_locator(AutoMinorLocator())
        #ax.xaxis.set_major_formatter(FormatStrFormatter("%d")) #ScalarFormatter())
        ax.spines["right"].set_visible(False)
        ax.spines["top"].set_visible(False)
        ax.grid(which='major', axis='y', linestyle='--')

        self.lns.append(ln)

        print(self.idx_it, self.it, self.mol)
        if self.project_type == 'water' and self.mol == 'H2O':
            self.dG_base_df = dg

    #TODO: merge with plot_dg
    def plot_dg_inhib(self, input_df):

        ax = self.ax_dG
        ax2 = self.ax2
        mol = self.mol
        mol2 = self.mol2
        lns = self.lns

        dg2, dg2_m, dg2_p = self.getDeltaG(input_df, self.it, self.idx_it)
        loc1, loc2 = self.loc1, self.loc2
        ranges = dg2.iloc[loc1:loc2].index.values
        ln2 = ax2.plot(ranges, dg2.iloc[loc1:loc2],'.', fillstyle='none', color=self.color2, label=f'5mM {mol2[mol]}')
        ax2.fill_between(ranges, dg2_m.iloc[loc1:loc2], dg2_p.iloc[loc1:loc2], alpha=0.2, color=self.color2)
        ax2.set_xlabel(r'$\itd$$_5$ ($\AA$)') #, color=self.color2)
        ax2.set_xlim(1.5, ax.get_xlim()[1])
        ax.set_xlim(1.5, ax.get_xlim()[1])
        ax2.set_xscale('log')
        #ax2.tick_params(axis='x', labelcolor=self.color2)
        ax2.set_ylim(-4, self.ax_dG.get_ylim()[1])
        lns.append(ln2)
        labs = [p[0].get_label() for p in lns]
        self.ax_dG.legend([ln[0] for ln in lns], labs,loc='upper right')
        #ax.set_title(mol)

        #ax2.set_ylim(-4, 6)

    def plot_regions(self, df):

        ax = self.ax_regions
        steps = self.steps
        count_total = df.size
        state_sampling_fraction = self.getStateSamplingFraction(df)

        steps_to_plot = [self.ax_dG.get_xlim()[0]]
        for i in steps:

            steps_to_plot.append(i)
        steps_to_plot.append(self.ax_dG.get_xlim()[1])

        sampling_fraction = [c for c in state_sampling_fraction]
        sampling_fraction.append(1.0)

        ax.step(steps_to_plot, sampling_fraction, where='post', linestyle=self.linestyle, color=self.color_mol[self.idx_it], label=self.it)

        if self.project_type == 'normal':
            ax.set_yscale('log')
            ax.set_yticks([1e-4, 1e-2, 1])

        ax.set_xlim(self.ax_dG.get_xlim())
        ax.set_xscale('log')
        ax.yaxis.set_minor_locator(AutoMinorLocator())
        ax.spines["right"].set_visible(False)
        ax.spines["top"].set_visible(False)

        if self.mol == 'ViAc':
            ax.set_xlabel(r'$\itd$$_{NAC,donor}$ ($\AA$)')
        elif self.project_type == 'water':
           ax.set_xlabel(r'$\itd$$_{NAC,hydrolysis}$ ($\AA$)')
           ax.set_ylim(0.8,1.05)

        else:
            ax.set_xlabel(r'$\itd$$_{NAC,acceptor}$ ($\AA$)')

        if (self.idx_mol == 0 and self.project_type == 'inhibition') or (self.idx_mol == 0 and self.mol == 'H2O'):
            ax.set_ylabel('Sampling\nfraction')

        if self.idx_it == len(self.iterables)-1:

            count_regions = self.count_regions[self.idx_it]
            label_ticks = []
            for i in self.steps:
                label_ticks.append(i)
            label_ticks.append(self.ax_regions.axes.get_xlim()[1])
            label_centers = np.mean(np.array([self.steps, label_ticks[1:]]), axis=0)

            if self.project_type == 'normal':
                label_location = 5e-4
            elif self.project_type == 'inhibition':
                label_location = 1e-1
            else:
                label_location = 0.85

            for idx_labels, x in enumerate(label_ticks): #
                if idx_labels >= 0:
                    ax.text(x, label_location, self.labels_regions[idx_labels], ha='right')
                if idx_labels < len(label_ticks)-1:
                    ax.vlines(x, 0, count_regions[idx_labels], linestyle='dotted', color='grey')

    def plot_combinatorial(self, df, opt_color=None):

        if isinstance(opt_color, np.ndarray):
            color = opt_color
        else:
            color = self.color_mol
        ax=self.ax_comb
        states = self.states

        #color = self.color_mol
        if self.project_type != 'normal':
            df.plot.barh(ax=ax,
                          color=color,
                          label=False),
                          #width=Plots.p_type[self.project_type]['bar_width'])
            #ax.set_xscale('log')
            ax.set_yticks(list(states.keys()))
            ax.set_xscale('log')
            ax.invert_yaxis()
            ax.set_xlabel('Sampling fraction')
            print(self.idx_mol)

            #plt.setp(ax.get_xtick(), visible=True)
            if self.idx_mol == 2:

                plt.setp(ax.get_yticklabels(), visible=False)
            else:
                ax.set_yticklabels(list(states.values())) #, rotation=0)
                ax.set_ylabel('States')
        else:
            df.plot.bar(ax=ax,
                        color=color,
                        label=False,
                        width=Plots.p_type[self.project_type]['bar_width'])

            ax.set_yscale('log')
            ax.set_xticks(list(states.keys()))
            ax.set_yticks([1e-6, 1e-4, 1e-2, 1])
            ax.set_xticklabels(list(states.values()), rotation=0)

            if self.project_type == 'normal':
                if self.idx_mol == 1:
                    ax.set_ylabel('Sampling fraction')
                    plt.setp(ax.get_xticklabels(), visible=False)
                elif self.idx_mol == len(self.projects)-1:
                    ax.set_xlabel('States')
                    ax.set_xticklabels(list(states.values()), rotation=0)

                else:
                    plt.setp(ax.get_xticklabels(), visible=False)
            elif self.project_type == 'inhibition':
                ax.set_ylabel('Sampling fraction')
                ax.set_xlabel('States')

        ax.grid(linestyle='--', axis='y')

        ax.legend().remove()

    def plot_wrapper(self,
                     dG_df : pd.DataFrame,
                     comb_df : pd.DataFrame):

        idx = self.idx_it
        list_its_ = comb_df.columns.get_level_values('l3').unique().to_list()

        if self.project_type == 'inhibition':
            list_its_.remove('5mM')

        comb_it_df = comb_df.loc[:, comb_df.columns.get_level_values('l3') == list_its_[idx]]

        self.counts_states[idx] =  np.asarray([np.count_nonzero(comb_it_df == i) for i in self.states]) / comb_it_df.size
        self.count_regions[idx] = self.getStateSamplingFraction(comb_it_df)

        self.plot_dg(dG_df)
        self.plot_regions(comb_it_df)

    def plot_metrics(self, correlation):

        #correlation = correlation_table[0]

        ax = self.ax_metric
        cmap= plt.get_cmap('YlGn') #'RdYlGn')
        cmap.set_bad('black')

        (state_labels, states) = list(self.states.values()), list(self.states.keys()),
        #cmap.set_bad('grey')

        if self.project_type == 'inhibition':
            label1, label2 = self.mol2[self.mol], self.mol
            ticks1, ticks2 = self.states_mol, states
            ticklabels1, ticklabels2 = self.state_labels_mol, state_labels
            vmin, _vmax = 0, 0.006156262816740886

        else:

            label2, label1 = 'water', self.mol
            ticks2, ticks1 = states, list(self.states_mol)
            ticklabels2, ticklabels1 = state_labels, list(self.state_labels_mol)
            correlation = correlation.T
            vmin, _vmax = 0, 0.008432989270410602

        print(ticks1, ticks2)
        print(ticklabels1, ticklabels2)

        vmin, vmax = 0, 1
        print(self.mol2[self.mol], self.mol, np.shape(correlation))
        corr_img = ax.imshow(correlation / _vmax, cmap=cmap, vmin=vmin, vmax=vmax, aspect='auto')

        ax.set_xticks(ticks1)
        ax.set_yticks(ticks2)
        ax.set_xticklabels(ticklabels1, rotation=90)
        ax.set_yticklabels(ticklabels2)

        ax.invert_yaxis()
        if self.project_type == 'inhibition':
            ax.set_xlabel(label1)
            ax.set_ylabel(label2)
            #ax.set_xticks(ticks1)
# =============================================================================
#
#             if self.idx_mol == len(self.projects)-1 :
#                 cbar = plt.colorbar(corr_img, ax=ax)
#                 cbar.set_label(self.metric_method)
# =============================================================================

        elif self.project_type == 'water':

            if self.it.split('_')[-1] == '5mM':
                it = f'{self.it.split("_")[0]}\n+\n{" ".join(self.it.split("_")[1:])}'
            else:
                it = self.it

            ax.set_ylabel(it, rotation='horizontal', va='center', ha='right')

            if self.idx_it != len(self.iterables)-1:
                plt.setp(ax.get_xticklabels(), visible=False)

            else:
                ax.set_xlabel(label1)

            if self.idx_mol == 1:
                plt.setp(ax.get_yticklabels(), visible=False)

        return corr_img

    def setSubplots(self, grid_dG, grid_comb):

        idx_mol = self.idx_mol
        project_type = self.project_type

        if idx_mol == 0:
            if project_type == 'normal':
                self.ax_dG = plt.subplot(grid_dG[0, idx_mol])
                self.ax_comb = plt.subplot(grid_comb[idx_mol,0])
                self.ax_regions = plt.subplot(grid_dG[1, idx_mol])
            elif project_type == 'inhibition':
                self.ax_dG = plt.subplot(grid_dG[0, idx_mol])
                self.ax_regions = plt.subplot(grid_dG[1, idx_mol])
            else:
                self.ax_dG = plt.subplot(grid_dG[0, 0])
                #self.ax_comb = plt.subplot(grid_comb[0,idx_mol])
                self.ax_regions = plt.subplot(grid_dG[1, 0])
        else:
            #ax = plt.subplot(grid_dG[0, idx_mol], sharey=ax)
            if project_type == 'normal':
                self.ax_dG = plt.subplot(grid_dG[0, idx_mol], sharey=self.ax_dG)
                self.ax_comb = plt.subplot(grid_comb[idx_mol,0], sharex=self.ax_comb, sharey=self.ax_comb)
            elif project_type == 'inhibition':
                self.ax_dG = plt.subplot(grid_dG[0, idx_mol], sharey=self.ax_dG)
                #self.ax_comb = plt.subplot(grid_comb[0, idx_mol], sharey=self.ax_comb)
            elif project_type == 'water':
                self.ax_comb = plt.subplot(grid_comb[0, idx_mol])

            self.ax_regions = plt.subplot(grid_dG[1, idx_mol], sharey=self.ax_regions)

        if project_type == 'inhibition':
            if self.metrics:
                sub_grid = gridspec.GridSpecFromSubplotSpec(2,2, subplot_spec = grid_comb[0, idx_mol], hspace=0, wspace=0, height_ratios=[4,10], width_ratios=[10,4])
                self.ax_metric = plt.subplot(sub_grid[1,0])
                self.ax_comb = plt.subplot(sub_grid[1,1], sharey=self.ax_metric)
                self.ax_comb2 = plt.subplot(sub_grid[0,0], sharex=self.ax_metric)

            else:
                self.ax_comb = plt.subplot(grid_comb[0,idx_mol])
                self.ax_comb2 = plt.subplot(grid_comb[0,idx_mol])

    def plotBindingProfileAndCombinatorial(self):
        figures = {}
        for project_type, projects in self.supra_project.items():
            print(project_type)
            input_states = self.input_states
            p_type = self.p_type
            dG = self.dG
            combinatorial = self.combinatorial
            plot_specs = self.plot_specs
            loc1, loc2 = self.loc1, self.loc2
            mol2 = self.mol2
            color2 = self.color2

            fig=plt.figure(figsize=p_type[project_type]['figsize'], constrained_layout=True) #
            (nrows, ncols) = p_type[project_type]['grid_layout']
            outer_grid = gridspec.GridSpec(nrows, ncols, hspace=0.3, height_ratios = p_type[project_type]['outer_grid_heights']) #, width_ratios = [3,1])

            if project_type == 'normal':
                grid_dG = gridspec.GridSpecFromSubplotSpec(2, len(projects), subplot_spec = outer_grid[0,:], hspace=0.3, wspace=0.1, height_ratios=[4,2])
                grid_comb = gridspec.GridSpecFromSubplotSpec(len(projects), 1, subplot_spec = outer_grid[2,:], hspace=0.2) #, wspace=0.1) #wspace=0.2)
                grid_figure = gridspec.GridSpecFromSubplotSpec(1, len(projects), subplot_spec = outer_grid[1,:], wspace=0.1)
            elif project_type == 'inhibition':
                grid_dG = gridspec.GridSpecFromSubplotSpec(2, len(projects), subplot_spec = outer_grid[0,:], hspace=0.3, wspace=0.1, height_ratios=[4,2])
                grid_comb = gridspec.GridSpecFromSubplotSpec(1, len(projects), subplot_spec = outer_grid[1,:], hspace=0.2, wspace=0.1)
                #grid_figure = gridspec.GridSpecFromSubplotSpec(1, 2, subplot_spec = outer_grid[2,:], wspace=0.1)
            elif project_type == 'water':
                grid_dG = gridspec.GridSpecFromSubplotSpec(2, 3, subplot_spec = outer_grid[0,:], height_ratios=[2,1], hspace=0.2, wspace=0.2)
                grid_comb = gridspec.GridSpecFromSubplotSpec(1,2, subplot_spec = outer_grid[1,1:], wspace=0.2)
                grid_figure = gridspec.GridSpecFromSubplotSpec(2,2, subplot_spec = outer_grid[2,1:], wspace=0.3, height_ratios=[10,2])
                #grid_corr = gridspec.GridSpecFromSubplotSpec(1, 1, subplot_spec = outer_grid[2,:], wspace=0.1)

            supra_df = pd.DataFrame()

            self.projects : dict = projects
            self.project_type : str = project_type
            self.labels_regions = list(self.subset_states[project_type].keys())

            for idx_mol, (mol, project) in enumerate(projects.items()):
                print(mol)

                if project_type != 'water':
                    dG_df = dG[project_type][mol]
                else:
                    base_df = dG[project_type][mol]
                    dG_df = base_df.iloc[:, base_df.columns.get_level_values(0) == mol].droplevel(0, axis=1)

                #TODOTODOTODOTODOTODOTODOTODOT for inhib and water. make function maybe
                combinatorial_df = combinatorial[project_type][mol]
                _states = np.unique(combinatorial_df.values.flat)
                self.states = {idx : input_states[project_type][k] for idx, k in enumerate(_states)}
                print(self.states)

                combinatorial_df.replace(to_replace=_states, value=self.states.keys())

                self.iterables = project.parameter
                self.steps = plot_specs[project_type][mol][0]
                self.color_mol = plot_specs[project_type][mol][1]
                self.linestyle = plot_specs[project_type]['linestyle']

                self.idx_mol = idx_mol
                self.mol = mol

                self.setSubplots(grid_dG, grid_comb)

                self.lns = []

                #ax_comb = plt.subplot(grid_comb[idx_mol])
                self.counts_states = np.empty([len(self.iterables), len(self.states)])
                self.count_regions = np.empty([len(self.iterables), len(self.steps)+1])

                for idx, it in enumerate(self.iterables):
                    self.it : str = it
                    self.idx_it = idx
                    self.plot_wrapper(dG_df, combinatorial_df)

                mol_df = pd.DataFrame(self.counts_states.T, index=list(self.states.keys()), columns=self.iterables)
                #mol_df = pd.DataFrame(counts_states, columns=list(states.keys()), index=iterables)
                supra_df = pd.concat([supra_df, mol_df], axis=1)
                if project_type != 'water':
                    self.plot_combinatorial(mol_df)

                #NOTE: loading of mol1 and mol1 + mol2 is done above. Here is only to handle mol2
                if project_type == 'inhibition':
                    print(mol)

                    self.ax2 = self.ax_dG.twiny()
                    self.color_mol = [color2]
                    self.linestyle = 'dashed'
                    self.idx_it = 0
                    self.it = '5mM'
                    inib_mol = mol2[mol]

                    inib_df = pd.read_csv(f'{project.results}/binding_profile/binding_profile_acylSer-His_{inib_mol}_mean_b0_e-1_s1.csv', index_col=0)
                    df_comb_inhib = combinatorial_df.loc[:, combinatorial_df.columns.get_level_values('l3') == self.it]

                    #TODO: fix dims
                    _states_mol = np.unique(df_comb_inhib.values.flat)
                    self.state_labels_mol = [v for k, v in self.states.items() if k in self.states_mol]

                    if self.metrics:
                        self.plot_metrics(self.discretized_matrixes[project_type][mol])
                    self.plot_dg_inhib(inib_df)
                    self.plot_regions(df_comb_inhib)
                    state_sampling_fraction = self.getStateSamplingFraction(df_comb_inhib, states= self.states_mol, w_subset=False)
                    x_values = self.ax_comb.get_xticks().flatten() + 0.35

                    #plot_combinatorial
                    print(inib_mol, self.states_mol, state_sampling_fraction)
                    self.ax_comb2.bar(self.states_mol, state_sampling_fraction, color=color2)
                    self.ax_comb2.set_yscale('log')

            #NOTE: This is only loading the other _water dfs. The df of water has been handled before
            if project_type == 'water':

                sub_w_grid_dG = gridspec.GridSpecFromSubplotSpec(2,2, subplot_spec = grid_dG[:,1:], height_ratios=[2,1], hspace=0.2, wspace=0.1)

                for idx_mol_w, (mol_w, iterables_mol_w2) in enumerate(self.mol_water.items()):
                    print(mol_w, iterables_mol_w2)

                        #self.plot_correlations(self.discretized_trajectories[project_type][mol_w])

                    base_df_mol = dG[project_type]['H2O']
                    dG_mol_df = base_df_mol.iloc[:, base_df_mol.columns.get_level_values(0) == mol_w].droplevel(0, axis=1)

                    color_mol = plot_specs['normal'][mol_w][1]
                    color_mol = np.vstack([color_mol, pl.cm.Oranges(np.linspace(0.5,0.5,1))])
                    color_mol_comb = np.vstack([pl.cm.Greys(np.linspace(1,1,1)), color_mol, pl.cm.Oranges(np.linspace(0.5,0.5,1))])

                    self.color_mol = color_mol
                    self.mol = mol_w
                    self.idx_mol = idx_mol_w

                    self.ax_dG = plt.subplot(sub_w_grid_dG[0, idx_mol_w])
                    self.ax_regions = plt.subplot(sub_w_grid_dG[1, idx_mol_w], sharey=self.ax_regions)
                    self.iterables = iterables_mol_w2[0] + [f'100mM_{iterables_mol_w2[1]}_5mM']
                    self.counts_states = np.empty([len(self.iterables), len(self.states)])
                    self.count_regions = np.empty([len(self.iterables), len(self.steps)+1])
                    if idx_mol_w == 0:
                        self.ax_comb = plt.subplot(grid_comb[0, idx_mol_w])
                    else:
                        self.ax_comb = plt.subplot(grid_comb[0, idx_mol_w], sharey=self.ax_comb)

                    _base_states_mol = np.unique(combinatorial_df.values.flat) #input_states['normal']
                    base_states_mol = {idx : input_states[project_type][k] for idx, k in enumerate(_base_states_mol)}
                    comb_df = combinatorial[project_type][mol]
                    comb_df = comb_df.iloc[:, comb_df.columns.get_level_values('l2') == mol_w]

                    _states = np.unique(comb_df.values.flat)

                    self.states = {idx : input_states[project_type][k] for idx, k in enumerate(_states)}
                    print(self.states)
                    comb_df.replace(to_replace=_states, value=self.states.keys())
                    self.states_mol = base_states_mol.keys()
                    self.state_labels_mol = base_states_mol.values()

                    if self.metrics:
                        sub_w_grid_metrics = gridspec.GridSpecFromSubplotSpec(len(self.iterables),1, subplot_spec = grid_figure[0, idx_mol_w])
                    for idx_w, it_w in enumerate(self.iterables):
                        self.it = it_w
                        self.idx_it = idx_w

                        #WARNING! This is assuming mol2 is the last iterable! make auto
                        if idx_w < len(self.iterables) -1 :
                            self.linestyle = plot_specs[project_type]['linestyle']
                        else:
                            self.linestyle = 'dashed'
                        #combinatorial_mol_df = comb_df.iloc[:, comb_df.columns.get_level_values('l2') == mol_w]

                        self.plot_wrapper(dG_mol_df, comb_df)

                        if self.metrics:
                            self.ax_corr = plt.subplot(sub_w_grid_metrics[idx_w,0])

                            corr_image = self.plot_metrics(self.discretized_matrixes[project_type][mol_w][it_w])

                    w_count_states = np.vstack([mol_df.values.T, self.counts_states])
                    mol_water_df = pd.DataFrame(w_count_states.T, index=list(self.states.keys()), columns=['0 mM'] + self.iterables)
                    self.plot_combinatorial(mol_water_df, opt_color=color_mol_comb)

                if self.metrics :

                    _colorbar = plt.subplot(grid_figure[1,:])
                    _colorbar.axis('off')
                    cbar = plt.colorbar(corr_image, ax=_colorbar, orientation='horizontal')
                    cbar.set_label(self.metric_method)
    # =============================================================================
    #             else:
    #                 ax_figure = fig.add_subplot(grid_figure[idx_mol])
    #                 if mol == 'BuOH':
    #                     image_mol = plt.imread(f'{project.results}/figures/msAcT_BuOH_100mM_activeSite_ATESB.png')
    #                 elif mol =='BeOH':
    #                     image_mol = plt.imread(f'{project.results}/figures/msAcT_BeOH_100mM_activeSite_ATESB.png')
    #                 else:
    #                     image_mol = plt.imread(f'{project.results}/figures/msAcT_ViAc_50mM_activeSite_ATESB.png')
    #
    #                 ax_figure.imshow(image_mol)
    #                 ax_figure.axis('off')
    # =============================================================================

            #plot_combinatorial(supra_df)
            #fig.tight_layout()
            fig.show()
            figures[project_type] = fig
            fig.savefig(f'{project.results}/figures/BindingProfile_combinatorial_{project_type}.png', dpi=600, bbox_inches='tight')

        return figures

    def getDeltaG(self, input_df, it, idx):

        if self.project_type == 'inhibition':
            #print(input_df)
            print(it, self.mol2[self.mol])
            if it == '100mM_{self.mol2[self.mol]}_5mM' or it == self.mol2[self.mol]:
                it = '5mM'

        dg = input_df[f'$\Delta$G {it}']
        if not isinstance(dg, pd.Series):
            dg = dg.iloc[:, idx]
            dg_index = [i for i, x in enumerate(input_df.columns.get_loc(dg.name)) if x][idx]
        else:
            dg_index = input_df.columns.get_loc(dg.name)
        dg_m = input_df.iloc[:, dg_index+1]
        dg_p = input_df.iloc[:, dg_index+2]

        return dg, dg_m, dg_p

    def getStateSamplingFraction(self, df, states=None, w_subset=True) -> list:

        count_total = df.size

        if states == None:
            states = self.states

        state_sampling_fraction = []
        if w_subset:
            for idx_regions, (label_state, subset) in enumerate(self.subset_states[self.project_type].items()):
                state_sampling_fraction.append(np.asarray([np.count_nonzero(df == i) / count_total for i in states if i in subset]).sum())
        else:
            state_sampling_fraction = np.asarray([np.count_nonzero(df == i) / count_total for i in states])

        return state_sampling_fraction

    def deprecated_plot_StateSampling_vs_scalars(self, df, ax):
        if self.project_type == 'water':
            iterable_scalars = [0]
            #WARNING! this is assuming mol2 is the last one on the list. make auto
            for it in self.iterables[:-1]:
                try:
                    bulk_value=float(str(it).split('M')[0])
                    unit='M'
                except:
                    bulk_value=float(str(it).split('mM')[0])
                    unit='mM'
                iterable_scalars.append(bulk_value)
            print(iterable_scalars)
            for state_col in df.columns:
                ax.plot(iterable_scalars, df[state_col].values[:-1])

def plot_layout(parameters):
    """Function to optimize plot layout based on number of subplots generated."""

    length=len(parameters)

    if length < 4:
        rows=1
        columns=length
    else:
        rows=2

        if not length % 2 == 0:
            columns=int(np.rint((length/2)+0.1))

        else:
            columns=int(length/2)

    if not length % 2 == 0 and length > 4:

        fix_layout=True
    else:
        fix_layout=False

    return rows, columns, fix_layout

def plot_feature_histogram(feature_dict):
    """Function to plot the histogram of raw features."""

    rows, columns, fix_layout=plot_layout(parameters)
    fig,axes=plt.subplots(rows, columns, sharex=True, constrained_layout=True)

    plt.suptitle(f'Feature: {name}')

    for plot, feature in zip(axes.flat, features.items()):
        for f in feature[1]:
            if not os.path.exists(f):
                feature[1].remove(f)
                print(f'\tCould not find data for {f}')

        data=np.asarray(pyemma.coordinates.load(feature[1]))
        plot.hist(data.flatten()) #, label=feature[0])
        plot.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
        plot.set_xlabel(r'{} ($\AA$)'.format(name))
        plot.set_ylabel('Counts')
        #plot.set_title(feature[0])

        data=[]

        fix_layout

    plt.savefig(f'{results}/histogram_{name}.png', bbox_inches="tight", dpi=600)

    return plt.show()

def plot_MFPT(mfpt_df, scheme, feature, parameters, error, regions=None, labels=None):
    """Function to plot heatmap of MFPTS between all states."""

    # Taken from matplotlib documentation. Make images respond to changes in the norm of other images (e.g. via the
    # "edit axis, curves and images parameters" GUI on Qt), but be careful not to recurse infinitely!
    def update(changed_image):
        for im in images:
            if (changed_image.get_cmap() != im.get_cmap() or changed_image.get_clim() != im.get_clim()):
                im.set_cmap(changed_image.get_cmap())
                im.set_clim(changed_image.get_clim())

    images = []

    rows, columns, fix_layout=plot_layout(parameters)
    fig, axes = plt.subplots(rows, columns, constrained_layout=True, figsize=(9,6))
    fig.suptitle(f'Discretization: {scheme}\nFeature: {feature} (error tolerance {error:.1%})', fontsize=14)

    cmap_plot=plt.cm.get_cmap("gist_rainbow")
    #cmap_plot.set_under(color='red')

    #cmap_plot.set_over(color='yellow')
    cmap_plot.set_bad(color='white')

    vmins, vmaxs=[], []

    for plot, parameter in zip(axes.flat, parameters):

        means=base.MSM.mfpt_filter(mfpt_df, scheme, feature, parameter, error)

        try:
            if scheme == 'combinatorial':
                contour_plot=plot.pcolormesh(means, edgecolors='k', linewidths=1, cmap=cmap_plot)
                label_names=base.Functions.sampledStateLabels(regions, sampled_states=means.index.values, labels=labels)
                positions=np.arange(0.5, len(label_names)+0.5)
                plot.set_xticks(positions)
                plot.set_xticklabels(label_names, fontsize=7, rotation=70)
                plot.set_yticks(positions)
                plot.set_yticklabels(label_names, fontsize=7)
            else:
                contour_plot=plot.pcolormesh(means, cmap=cmap_plot)
                ticks=plot.get_xticks()+0.5
                plot.set_xticks(ticks)
                plot.set_xticklabels((ticks+0.5).astype(int))
                plot.set_yticks(ticks[:-1])
                plot.set_yticklabels((ticks+0.5).astype(int)[:-1])

            plot.set_facecolor('white')
            plot.set_title(parameter) #, fontsize=8)
            plot.set_xlabel('From state', fontsize=10)
            plot.set_ylabel('To state', fontsize=10)
            images.append(contour_plot)

        except:
            print('No values to plot')

    # Find the min and max of all colors for use in setting the color scale.
    vmins=[]
    vmaxs=[]
    for image in images:
        array=image.get_array()
        try:
            vmin_i=np.min(array[np.nonzero(array)])
        except:
            vmin_i=1
        try:
            vmax_i=np.max(array[np.nonzero(array)])
        except:
            vmax_i=1e12
        vmins.append(vmin_i)
        vmaxs.append(vmax_i)

    vmin=min(vmins)
    vmax=max(vmaxs)
    #vmax = max(image.get_array().max() for image in images)

    norm = colors.LogNorm(vmin=vmin, vmax=vmax)

    for im in images:
        im.set_norm(norm)

    print(f'limits: {vmin:e}, {vmax:e}')

    cbar=fig.colorbar(images[-1], ax=axes)
    cbar.set_label(label=r'MFPT (ps)', size='large')
    cbar.ax.tick_params(labelsize=12)
    for im in images:
        im.callbacksSM.connect('changed', update)

    return images

def plot_flux(flux_df, ligand, results):
    """Function to plot fluxes from a dataframe of fluxes."""

    schemes=flux_df.index.unique(level='Scheme')
    features=flux_df.index.unique(level='feature')

    for scheme in schemes:
        for feature in features:

            properties=flux_df.columns

            fig, axes = plt.subplots(1,len(properties), sharex=True, constrained_layout=True)
            fig.suptitle(f'Discretization: {scheme}\n Feature: {feature}')

            for ax, prop in zip(axes, properties):

                flux_df.loc[(scheme, feature), prop].plot(linestyle='-', marker='o', ax=ax, title=prop)
                ax.set_xlabel(f'[{ligand}] (M)')
                ax.set_ylabel(prop+r' ($s^{-1}$)')
            plt.savefig(f'{results}/netFlux_{scheme}-{feature}.png', dpi=600)
            plt.show()

def plot_pathways(pathway_df, ligand, results):
    """Function to plot pathways from a dataframe of pathways."""

    schemes=pathway_df.index.unique(level='Scheme')
    features=pathway_df.index.unique(level='feature')

    for scheme in schemes:
        for feature in features:

            print(scheme, feature)

            pathway_df.loc[(scheme, feature)].dropna(axis=1, how='all').plot(linestyle='-', marker='o')
            plt.title(f'Discretization: {scheme}\n Feature: {feature}')
            plt.xlabel(f'[{ligand}] (M)')
            plt.ylabel(r'Pathway flux ($s^{-1}$)')
            plt.savefig(f'{results}/pathways_{scheme}-{feature}.png', dpi=600)
            plt.show()

def plot_committor(committor_df, ligand, results, regions, labels):
    """Function to plot fluxes from a dataframe of fluxes."""

    schemes=committor_df.index.unique(level='Scheme')
    features=committor_df.index.unique(level='feature')
    parameters=committor_df.index.unique(level='parameter')

    committors=('Forward', 'Backward')

    for scheme in schemes:
        for feature in features:

            fig, axes = plt.subplots(1,len(committors), sharey=True, constrained_layout=True)
            for parameter in parameters:
                for ax, committor in zip(axes, committors):
                    try:
                        df_c=committor_df.loc[(scheme, feature, parameter), committor].dropna()

                        if scheme == 'combinatorial':

                            df_c.plot(linestyle='-', marker='o', ax=ax, label=parameter, title=committor)

                            sampled_states=committor_df.unstack(level=2).loc[(scheme, feature)][committor].dropna().index.unique(level='states').values
                            label_names=base.Functions.sampledStateLabels(regions, sampled_states=sampled_states, labels=labels)

                            ax.set_xticks(sampled_states)
                            ax.set_xticklabels(label_names, rotation=70)

                        else:
                            df_c.plot(ax=ax, label=parameter, title=committor)
                        ax.set_ylabel('Committor Probability')
                    except:
                        print('no plot at ', parameter)

            plt.suptitle(f'Discretization: {scheme}\n Feature: {feature}')

            plt.legend(title=f'[{ligand}] (M)')
            plt.savefig(f'{results}/committors_{scheme}-{feature}.png', dpi=600)
            plt.show()


__author__ = 'luisl_000'

# This class contains helper functions to create objects or animations

from components import Animator
from util_math import Vector2
from os import listdir
from re import split
from pygame import Surface
from pygame import image
from components import RigidBody
from components import BoxCollider

def set_lamp_light_attributes(lamp_light, rs):
    lamp_light.renderer.depth = 10000
    lamp_light.add_component(BoxCollider(50, 50))
    lamp_light.collider.is_trigger = True
    lamp_light.tag = "lamp light"
    rs.light_sources.append(lamp_light)

def set_floor_attributes(floor):
    floor.collider.restitution = 0
    floor.collider.surface_friction = 0.75
    floor.renderer.depth = -10
    floor.tag = "floor"

def set_wall_attributes(wall):
    wall.collider.restitution = 0
    wall.collider.surface_friction = 0.8
    wall.tag = "wall"

def set_ceiling_attributes(ceiling):
    ceiling.collider.restitution = 0
    ceiling.collider.surface_friction = 0.8
    ceiling.tag = "ceiling"

def set_platform_attributes(platform):
    platform.renderer.depth = 1
    platform.collider.restitution = 0
    platform.collider.surface_friction = 0.75
    platform.tag = "platform"

def set_box_attributes(box):
    box.renderer.depth = 2
    box.collider.restitution = 0
    box.collider.surface_friction = 0.8
    box.collider.box.w -= 10
    box.collider.box.h -= 10

    box.add_component(RigidBody())
    box.rigid_body.velocity = Vector2(0.0, 0.0)
    box.rigid_body.gravity_scale = 2.0
    box.tag = "box"

def get_files_in_dir(dir_path):

    directory = listdir(dir_path)

    # a list to store the paths to the individual files
    file_paths = list()
    for file_ in directory:

        # add the whole relative path
        file_paths.append(dir_path + file_)

    # do a natural sort on the file names, meaning that string numbers are sorted
    # by numeric values.
    file_paths = natural_sort(file_paths)
    return file_paths

# This goes to a directory where each file represents an individual
# animation frame. This returns an animation object with the loaded frames.
def load_anim_from_directory(dir_path):

    file_list = get_files_in_dir(dir_path)

    # set up animation
    animation = Animator.Animation()
    for file_ in file_list:
        frame = image.load(file_).convert_alpha()
        animation.add_frame(frame)

    return animation

# Excerpt from stack overflow, Mark Byers
def natural_sort(l):
    convert = lambda text: int(text) if text.isdigit() else text.lower()
    alphanum_key = lambda key: [convert(c) for c in split('([0-9]+)', key)]
    return sorted(l, key=alphanum_key)

# given a tile image, this function will create a quilt of tiles onto a another surface
# given the width and height.
def create_img_from_tile(tile_surface, width, height):

    tile_w = tile_surface.get_width()
    tile_h = tile_surface.get_height()

    size = (width, height)
    dst_surface = Surface(size).convert()

    # make the dst_surface transparent
    color_mask = (7, 13, 17)
    dst_surface.fill(color_mask)
    dst_surface.set_colorkey(color_mask)

    # fill the dst_surface horizontally first with the tile, once we get to the edge
    # go down a column and repeat.
    for y in range(0, height, tile_h):
        for x in range(0, width, tile_w):
            dst_surface.blit(tile_surface, (x, y))

    return dst_surface

# conjoins surface b on the bottom of surface a
def conjoin_surfaces_vertically(surface_a, surface_b):
    width = max(surface_a.get_width(), surface_b.get_width())
    height = surface_a.get_height() + surface_b.get_height()

    size = (width, height)
    dst_surface = Surface(size).convert()

    # make the dst_surface transparent
    color_mask = (123, 54, 33)
    dst_surface.fill(color_mask)
    dst_surface.set_colorkey(color_mask)

    # center each surface relative to the destination surface
    center_a = width/2 - surface_a.get_width()/2
    center_b = width/2 - surface_b.get_width()/2

    dst_surface.blit(surface_a, (center_a, 0))
    dst_surface.blit(surface_b, (center_b, surface_a.get_height()))

    return dst_surface


import abc
from unittest import TestCase

import numpy as np

from emissor.representation.annotation import Entity
from emissor.representation.scenario import Annotation, class_type, module_source, class_source

class TestScenarioModule(TestCase):
    def test_annotation_from_object_with_module_source(self):
        annotation_value = Entity("uri:test")
        annotation = Annotation(class_type(annotation_value), annotation_value, module_source(__name__), 0)

        self.assertEqual("python-type:emissor.representation.annotation.Entity", annotation.type)
        self.assertEqual("python-source:test_emissor.test_representation.test_scenario", annotation.source)

    def test_annotation_from_object_with_class_source(self):
        annotation_value = Entity("uri:test")
        annotation = Annotation(class_type(annotation_value), annotation_value, class_source(Annotation), 0)

        self.assertEqual("python-type:emissor.representation.annotation.Entity", annotation.type)
        self.assertEqual("python-source:emissor.representation.scenario", annotation.source)

    def test_annotation_from_object_with_class_source_type(self):
        annotation_value = Entity("uri:test")
        annotation = Annotation(class_type(annotation_value), annotation_value, class_source(Annotation, include_type=True), 0)

        self.assertEqual("python-type:emissor.representation.annotation.Entity", annotation.type)
        self.assertEqual("python-source:emissor.representation.scenario.Annotation", annotation.source)

    def test_class_type_from_metaclass(self):
        """Test that class type is not resolved to abc.ABCMeta"""
        class TestABC(abc.ABC):
            pass

        class TestClass(TestABC):
            def __init__(self, name):
                self.name = name

            @classmethod
            def factory(cls):
                return cls(class_type(cls))

        annotation_value = TestClass.factory()
        annotation = Annotation(class_type(annotation_value), annotation_value, module_source(__name__), 0)

        self.assertRegexpMatches(annotation.type, ".*[.]TestScenarioModule.test_class_type_from_metaclass.<locals>.TestClass")
        self.assertEqual(annotation_value.name, annotation.type)
        self.assertEqual(annotation_value.name, class_type(TestClass))

    def test_annotation_from_numpy(self):
        annotation_value = np.array([])
        annotation = Annotation(class_type(annotation_value), annotation_value, module_source(np.__name__), 0)

        self.assertEqual("python-type:numpy.ndarray", annotation.type)
        self.assertRegexpMatches(annotation.source, "python-source:numpy#\d[.][\w.]+")


import argparse
import os
import pickle
import sys
import cv2
import numpy as np
import open3d as o3d
import matplotlib.pyplot as plt
from gta_utils import LIMBS, read_depthmap
from sklearn.neighbors import NearestNeighbors
from tqdm import tqdm       # some fancy visual effect
import glob
from utils import *
import copy
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--data_root', type=str, default='/GTA-1M/FPS-30')
parser.add_argument('--sequence_id', type=str, default='2020-05-20-21-13-13_resize')
parser.add_argument('--save_root', type=str, default='/GTA-1M/FPS-30/preprocessed_data')

# output 2D proximity maps and many other things in /GTA-1M/FPS-30/preprocessed_data

args = parser.parse_args()

if __name__ == '__main__':
    ##################### set up parameters #########################
    MAX_DEPTH = 20.0
    h, w = 256, 448
    # h, w = 1080, 1920
    scale = 1080/h
    fps = 30
    depth_inpaint_itr = 500

    ##################### join paths #########################
    save_feature_img_path = os.path.join(args.save_root, args.sequence_id, 'bps_feature_img')
    save_feature_npy_path = os.path.join(args.save_root, args.sequence_id, 'bps_feature_npy')
    save_depth_img_path = os.path.join(args.save_root, args.sequence_id, 'depth_inpaint_img')
    save_depth_npy_path = os.path.join(args.save_root, args.sequence_id, 'depth_inpaint_npy')
    save_rgb_path = os.path.join(args.save_root, args.sequence_id, 'rgb_img_input')

    info = pickle.load(open(os.path.join(args.data_root, args.sequence_id, 'info_frames.pickle'), 'rb'))
    info_npz = np.load(os.path.join(args.data_root, args.sequence_id, 'info_frames.npz'))

    ##################### create directories #########################
    if not os.path.exists(save_feature_img_path):
        os.makedirs(save_feature_img_path)
    if not os.path.exists(save_feature_npy_path):
        os.makedirs(save_feature_npy_path)
    if not os.path.exists(save_depth_img_path):
        os.makedirs(save_depth_img_path)
    if not os.path.exists(save_depth_npy_path):
        os.makedirs(save_depth_npy_path)
    if not os.path.exists(save_rgb_path):
        os.makedirs(save_rgb_path)

    # base images
    rgb_list = glob.glob(os.path.join(args.data_root, args.sequence_id, '*.jpg'))       # collect paths of all .jpg images under the directory
    rgb_list.sort()
    n_frame = int(rgb_list[-1][-9:-4]) + 1  # total frame number

    vis = o3d.visualization.Visualizer()
    vis.create_window(width=w, height=h, visible=True)

    ##################### construct point clouds for a sequence of frames #########################
    seq_cnt = 0
    start_frame = fps - 1       # skip the first fsp-1 frames
    end_frame = n_frame - fps*4     # skip the last fps*4 frames
    # todo: step = fps*1 if enlarge training size
    for frame_N in tqdm(range(start_frame, end_frame, fps*3)):    # N_th frame: to get inpainted depth map --> reference frame
        global_pcd = o3d.geometry.PointCloud()

        start = max(0, frame_N-fps*6) if fps==5 else max(5, frame_N-fps*6)      # include the past fps*6 frames
        end = frame_N + fps*4     # include the future fps*4 frames
        step = 3 * fps // 5     # step (@Di: don't know why this step is used)

        # @Benjamin: We go for a range of a number of frames centered around frame_N - fps*6 to frame_N + fps*4 :O
        # I don't understand this. E.g. frame_N = 500, fps=30 gives frame_cnt: 320, 338, 356, ..., 608
        # Oh I think it's about getting a local point cloud for the current short time period.
        # @Di: Agree.
        for frame_cnt in range(start, end, step):   # 10s
            img_path = os.path.join(args.data_root, args.sequence_id, '{:05d}'.format(frame_cnt) + '.jpg')
            depth_path = os.path.join(args.data_root, args.sequence_id, '{:05d}'.format(frame_cnt) + '.png')
            human_mask_path = os.path.join(args.data_root, args.sequence_id, '{:05d}'.format(frame_cnt) + '_id.png')

            ##################### read and clean data #########################
            img = cv2.imread(img_path)      # read RGB image [h,w,3]
            img = img[:, :, ::-1]  # BGR --> RGB (@Di: don't know how this works)
            img = cv2.resize(img, (w, h), interpolation=cv2.INTER_NEAREST)

            # @Di: why do we need this parameter?
            infot = info[frame_cnt]
            cam_near_clip = infot['cam_near_clip']  # near/far: distances from the camera to start/stop rendering.
            if 'cam_far_clip' in infot.keys():
                cam_far_clip = infot['cam_far_clip']
            else:
                cam_far_clip = 800.
            depth = read_depthmap(depth_path, cam_near_clip, cam_far_clip)  # read depth image [1080, 1920, 1], in meters
            depth = cv2.resize(depth, (w, h), interpolation=cv2.INTER_NEAREST)   # [h,w]
            # plt.imshow(depth, cmap='plasma')
            # plt.show()

            # obtain the human mask
            joints_2d = info_npz['joints_2d']
            p = joints_2d[frame_cnt, 0]   # (x,y) coordinate of head joint
            id_map = cv2.imread(human_mask_path, cv2.IMREAD_ANYDEPTH)   # read object & human mask [1080, 1920]
            human_id = id_map[np.clip(int(p[1]), 0, 1079), np.clip(int(p[0]), 0, 1919)]   # scalar, uint8, human id
            # @Benjamin: If head is not in image?
            # @Di: This is just to extract the RGB value of human being in the image. np.clip() ensures that this will work
            # even if the human head is not in the image.
            mask_human = id_map == human_id     # extract human mask from object & human mask [1080, 1920], true/false, 1:person, 0:background
            mask_human = cv2.resize(mask_human.astype(np.uint8) , (w, h))
            kernel = np.ones((3, 3), np.uint8)
            mask_dilation = cv2.dilate(mask_human.astype(np.uint8), kernel, iterations=1)       # expand the mask a bit

            # @Di: why?
            depth[depth > MAX_DEPTH] = 0    # exclude points with large depth value
            depth[mask_dilation==True] = 0    # exclude person in the depth map

            #################### reconstruct point cloud from consecutive frames #######################
            cam_int = info_npz['intrinsics'][frame_cnt]  # camera intrinsic params [3,3]
            # retrieve focal lengths
            focal_length_x = cam_int[0, 0]
            focal_length_y = cam_int[1, 1]

            depth = np.expand_dims(depth, axis=-1)      # @Di: why is axis -1? [h,w,1]

            # from rgb image and human excluded depth map
            rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
                o3d.geometry.Image(img),
                o3d.geometry.Image(depth.astype(np.float32)),
                depth_scale=1.0,
                depth_trunc=MAX_DEPTH,
                convert_rgb_to_intensity=False,
            )
            # create coloured point cloud from rgbd image
            pcd = o3d.geometry.PointCloud.create_from_rgbd_image(
                rgbd_image,
                o3d.camera.PinholeCameraIntrinsic(w, h, focal_length_x/scale, focal_length_y/scale, (w/2), (h/2)),
            )  # cam coordinate

            depth_pts = np.asarray(pcd.points)     # [1920*1080-n_human_mask, 3], coordinate of each pixel in the depth map
            depth_pts_aug = np.hstack([depth_pts, np.ones([depth_pts.shape[0], 1])])        # euclidean coordinates to homogeneous coordinates
            cam_extr_ref = np.linalg.inv(info_npz['world2cam_trans'][frame_cnt])        # camera extrinsic params
            depth_pts = depth_pts_aug.dot(cam_extr_ref)[:, :3]      # camera coordinates to world coordinates
            pcd.points = o3d.utility.Vector3dVector(depth_pts)

            # we put the point cloud into the global point cloud
            global_pcd.points.extend(pcd.points)
            global_pcd.colors.extend(pcd.colors)

        # the global point cloud has all information (including occluded points behind human for the current frame)
        # therefore we extract a new current frame depth image from global point cloud to use instead of input data
        # depth image.
        ##################### capture depth/RGB image of cur_frame_N from point cloud ###############
        for cur_frame_N in range(frame_N-5, frame_N+1):
            vis.add_geometry(global_pcd)
            cam_ext = np.transpose(info_npz['world2cam_trans'][cur_frame_N])
            cam_int = info_npz['intrinsics'][cur_frame_N]
            ctr = vis.get_view_control()
            cam_param = ctr.convert_to_pinhole_camera_parameters()
            cam_param.intrinsic= o3d.camera.PinholeCameraIntrinsic(w, h, cam_int[0,0]/scale, cam_int[1,1]/scale, w/2-0.5, h/2-0.5)
            cam_param.extrinsic = cam_ext
            ctr.convert_from_pinhole_camera_parameters(cam_param)
            vis.poll_events()
            vis.update_renderer()

            # better than input depth and rgb images for current frame
            out_depth = np.asarray(vis.capture_depth_float_buffer())
            # plt.imshow(out_depth)
            # plt.show()
            out_image = np.asarray(vis.capture_screen_float_buffer())
            # plt.imshow(out_image)
            # plt.show()

            vis.remove_geometry(global_pcd)

            ######################## inpaint/save depth map of cur_frame_N ##############################
            infot = info[cur_frame_N]
            cam_near_clip = infot['cam_near_clip']
            if 'cam_far_clip' in infot.keys():
                cam_far_clip = infot['cam_far_clip']
            else:
                cam_far_clip = 800.
            depth = read_depthmap(os.path.join(args.data_root, args.sequence_id, '{:05d}'.format(cur_frame_N) + '.png'), cam_near_clip, cam_far_clip)
            depth = cv2.resize(depth, (w, h), interpolation=cv2.INTER_NEAREST)  # [h,w]
            mask_sky_cur_frame_N = depth > MAX_DEPTH        # sky mask

            joints_2d = info_npz['joints_2d']
            p = joints_2d[cur_frame_N, 0]  # (x,y) coordinate of head joint
            # human & object mask
            id_map = cv2.imread(os.path.join(args.data_root, args.sequence_id, '{:05d}'.format(cur_frame_N) + '_id.png'), cv2.IMREAD_ANYDEPTH)  # [1080, 1920]
            human_id = id_map[np.clip(int(p[1]), 0, 1079), np.clip(int(p[0]), 0, 1919)]
            mask_human_cur_frame_N = id_map == human_id     # human mask (true/false)
            mask_human_cur_frame_N = cv2.resize(mask_human_cur_frame_N.astype(np.uint8), (w, h))

            # (out_depth == 0) * (mask_human_cur_frame_N == false): sky
            # (out_depth == 0): human or sky
            depth_mask_human = (out_depth == 0) ^ (out_depth == 0) * (mask_human_cur_frame_N == False) # human mask
            depth_mask_sky = (out_depth == 0) * (mask_human_cur_frame_N == False) # sky
            depth_mask_human_ratio = np.sum(depth_mask_human)/(h*w) # ratio of human to the image
            # print(depth_mask_ratio)
            if depth_mask_human_ratio <= 0.01:      # if ratio of human is small enough
                if depth_mask_human_ratio == 0:     # if no human at all
                    inpainted_depth_map = out_depth     # this is the inpainted depth map
                    inpainted_depth_mask_human_ratio = 0
                else:       # if the ratio is relatively large
                    print('interpolating depth for frame {}'.format(cur_frame_N))       # interpolate the depth map
                    # @Di: where is this function l1_inpainting() from?
                    inpainted_depth_map = l1_inpainting(out_depth,
                                                        depth_mask_human,
                                                        maxIter=depth_inpaint_itr)
                    # plt.imshow(inpainted_depth_map)
                    # plt.show()
                    # (out_depth > 0): not human or sky
                    inpainted_depth_mask_human = inpainted_depth_map[depth_mask_human] <= np.min(out_depth[out_depth>0])
                    inpainted_depth_mask_human_ratio = np.sum(inpainted_depth_mask_human) / (h*w)

                if inpainted_depth_mask_human_ratio <= 0.001:       # if ratio of human is small enough
                    inpainted_depth_map = inpainted_depth_map.astype(np.float32)        # take as the interpolated inpainted depth map

                    ################### get scene point cloud of cur_frame_N ###################
                    rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
                        o3d.geometry.Image((np.asarray(out_image)*255).astype(np.uint8)),
                        o3d.geometry.Image(np.asarray(inpainted_depth_map.astype(np.float32))),
                        depth_scale=1.0,
                        depth_trunc=MAX_DEPTH,
                        convert_rgb_to_intensity=False,
                    )
                    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(
                        rgbd_image,
                        o3d.camera.PinholeCameraIntrinsic(w, h, cam_int[0,0]/scale, cam_int[1,1]/scale, w/2-0.5, h/2-0.5),
                    )  # cam coordinate

                    scene_verts = np.asarray(pcd.points)     # [h*w, 3], coordinate of each pixel in the depth map
                    scene_verts_aug = np.hstack([scene_verts, np.ones([scene_verts.shape[0], 1])])
                    cam_extr_ref = np.linalg.inv(info_npz['world2cam_trans'][cur_frame_N])
                    scene_verts = scene_verts_aug.dot(cam_extr_ref)[:, :3]      # camera to world transformation
                    pcd.points = o3d.utility.Vector3dVector(scene_verts) # world coordinate

                    if scene_verts.shape[0] == np.sum(inpainted_depth_map > 0):     # if all points in the depth map has a corresponding point in the point cloud
                        ######################## save rgb/depth image ##############################
                        print('save data for cur_frame_N = {}'.format(cur_frame_N))
                        seq_cnt += 1        # for every cur_frame_N, if
                        plt.imsave('{}/seq_{:04d}_fr_{:05d}.png'.format(save_depth_img_path, seq_cnt, cur_frame_N), inpainted_depth_map)
                        np.save('{}/seq_{:04d}_fr_{:05d}.npy'.format(save_depth_npy_path, seq_cnt, cur_frame_N), inpainted_depth_map)

                        # save rgb img with person
                        img_cur_frame_N = cv2.imread(os.path.join(args.data_root, args.sequence_id, '{:05d}'.format(cur_frame_N) + '.jpg'))
                        img_cur_frame_N = img_cur_frame_N[:, :, ::-1]  # BGR --> RGB
                        img_cur_frame_N = cv2.resize(img_cur_frame_N, (w, h), interpolation=cv2.INTER_NEAREST)

                        rgb_image = copy.deepcopy(out_image)  # in [0,1]
                        rgb_image[mask_human_cur_frame_N == True] = img_cur_frame_N[mask_human_cur_frame_N == True] / 255.0 # scaled rgb value for human
                        rgb_image[mask_sky_cur_frame_N == True] = img_cur_frame_N[mask_sky_cur_frame_N == True] / 255.0 # scaled rgb value for sky
                        plt.imshow(rgb_image)
                        plt.imsave('{}/seq_{:04d}_fr_{:05d}.jpg'.format(save_rgb_path, seq_cnt, cur_frame_N), rgb_image)

                        ################### calculate body bps feature for all frames #######################
                        # some params (@Di: no idea why parameters are like this, maybe empirically determined)
                        start = cur_frame_N - fps//5 * 4
                        end = cur_frame_N + fps * 2 + 1
                        step = fps // 5
                        for cur_frame in range(start, end, step):
                            depth_flat = inpainted_depth_map.reshape(h*w)  # inpainted depth map
                            depth_mask_ind = np.where(depth_flat == 0)[0]  # indices of masked pixels (human or sky)
                            depth_nomask_ind = np.asarray(list(set(range(h*w))-set(depth_mask_ind)))  # indices of nonmasked pixels
                            depth_mask_sky_ind = ((out_depth == 0) * (mask_human_cur_frame_N==False)).reshape(h*w)  # indices of sky

                            # ############ distance to 21 joints ############
                            # nbrs = NearestNeighbors(n_neighbors=1, leaf_size=1, algorithm="ball_tree").fit(body_joints_3d)
                            # neigh_dist, neigh_ind = nbrs.kneighbors(scene_verts)
                            # body_bps = neigh_dist[:, 0]  # distance  [n_bps]  n_bps = h*w-n_human_mask
                            # print('max/min body bps feature value:', np.max(body_bps), np.min(body_bps))
                            # print(np.sum(body_bps))

                            ############ distance to skeleton (21 joints)  ############
                            # min distance from each pixel point P to a set of line segments AB
                            body_joints_3d = info_npz['joints_3d_world'][cur_frame]  # [21, 3] world coordinate
                            # 3D coordinates of the two end points of each limbs
                            A = body_joints_3d[np.asarray(LIMBS)[:, 0]]  # [n_limb, 3]
                            B = body_joints_3d[np.asarray(LIMBS)[:, 1]]  # [n_limb, 3]
                            n_pt = scene_verts.shape[0]
                            n_limb = A.shape[0]

                            A = np.tile(A, (n_pt, 1)).reshape(n_pt*n_limb, 3)  # [n_pt, n_limb, 3], n_pt=n_bps=scene_verts.shape[0]
                            B = np.tile(B, (n_pt, 1)).reshape(n_pt*n_limb, 3)
                            P = np.tile(scene_verts, n_limb).reshape(n_pt*n_limb, 3)

                            # compute distance from P to each limb
                            AB = B - A
                            AP = P - A
                            BP = P - B
                            temp_1 = np.multiply(AB, AP).sum(axis=-1)  # [n_pt, n_limb]
                            temp_2 = np.multiply(-AB, BP).sum(axis=-1)  # [n_pt, n_limb]
                            mask_1 = np.where(temp_1 <= 0)[0]   # angle between AB and AP >= 90
                            mask_2 = np.where((temp_1 > 0) * (temp_2 <= 0))[0]  # angle between AB and AP < 90 and angle between BA and BP >= 90
                            mask_3 = np.where((temp_1 > 0) * (temp_2 > 0))[0]   # angle between AB and AP < 90 and angle between BA and BP < 90
                            if len(mask_1) + len(mask_2) + len(mask_3) != n_pt*n_limb:
                                print('[distance calculation] num of verts does not match!')

                            dist_1 = np.sqrt(np.sum((P[mask_1]-A[mask_1])**2, axis=-1))  # [n_mask_1]
                            dist_2 = np.sqrt(np.sum((P[mask_2]-B[mask_2])**2, axis=-1))  # [n_mask_2]

                            x = np.multiply(AB[mask_3], AP[mask_3]).sum(axis=-1) / np.multiply(AB[mask_3], AB[mask_3]).sum(axis=-1)  # [n_mask_3]
                            x = x.repeat(3).reshape(-1,3)
                            C = x * AB[mask_3] + A[mask_3]  # C: [n_mask_3, 3], the projected point of P on line segment AB
                            dist_3 = np.sqrt(np.sum((P[mask_3]-C)**2, axis=-1))  # n_mask_3

                            dist = np.zeros(n_pt*n_limb)
                            dist[mask_1] = dist_1
                            dist[mask_2] = dist_2
                            dist[mask_3] = dist_3
                            dist = dist.reshape(n_pt, n_limb)  # [n_pt, n_limb], distance from each point in scene verts to each limb
                            body_bps = np.min(dist, axis=-1)   # [n_pt]
                            # print('max/min body bps feature value:', np.max(body_bps), np.min(body_bps))

                            ########### back to image plane, visualize bps feature map ###########
                            # print(cur_frame, cur_frame_N)
                            body_bps_full = np.zeros([h*w])
                            body_bps_full[depth_nomask_ind] = body_bps
                            body_bps_full[depth_mask_ind] = 0         # set masked pixels to 0
                            body_bps_full[depth_mask_sky_ind] = 100.0  # set sky pixels to 100
                            body_bps_full = body_bps_full.reshape((h, w))  # body bps feature map, [h, w]
                            body_bps_full = body_bps_full.astype(np.float32)
                            np.save('{}/seq_{:04d}_fr_{:05d}.npy'.format(save_feature_npy_path, seq_cnt, cur_frame), body_bps_full)

                            # plt.imshow(body_bps_full, cmap='plasma')
                            # plt.show()

                            body_bps_full[body_bps_full>10] = 10
                            fig = plt.imshow(body_bps_full, cmap='plasma', vmin=0, vmax=5)
                            plt.axis('off')
                            fig.axes.get_xaxis().set_visible(False)
                            fig.axes.get_yaxis().set_visible(False)
                            plt.imsave('{}/seq_{:04d}_fr_{:05d}.png'.format(save_feature_img_path, seq_cnt, cur_frame), body_bps_full, cmap='plasma')
                            # plt.show()

    print('total seq cnt:', seq_cnt)


#Задание 1
import random

def f1(numbers):
    # Создаем новый массив, в котором каждое число возводится в квадрат
    retnumbers = [num ** 2 for num in numbers]
    return retnumbers

# Создаем массив из 20 случайных чисел
numbers = []
for x in range(20):
    numbers.append(random.randint(1, 9))

# Выводим исходный массив на экран
print("Исходный массив:")
print(numbers)

# Возводим каждое число в квадрат
retnumbers = f1(numbers)

# Выводим новый массив на экран
print("Массив с числами, возведенными в квадрат:")
print(retnumbers)
print('###############################################')
#ЗАдание2

def q1(number):                                              #Функция q1 принимает число number и возвращает True
    return number % 5 != 0                                   #если число не делится на 5, и False в противном случае.

def q2(number):                                              #Функция q2 принимает число number и возвращает True
    return len(set(str(number))) == len(str(number))         #если число содержит уникальные цифры  и False в противном случае

def q3(numbers):                                             #Функция q3 принимает массив чисел numbers и применяет последовательно два фильтра
    # Убираем числа, кратные 5                                # сначала убирает числа, кратные 5, затем убирает повторы и
    numbers = list(filter(q1, numbers))                      #наконец, убирает числа с одинаковыми цифрами. Возвращается отфильтрованный массив чисел.
    # Убираем повторы
    numbers = list(set(numbers))
    # Убираем числа с одинаковыми цифрами
    numbers = filter(q2, numbers)
    return numbers
# Создаем массив из 100 случайных чисел
numbers = [random.randint(10, 100) for _ in range(100)]

# Применяем фильтры к массиву чисел
filternum = q3(numbers)

# Разбиваем отфильтрованные числа на группы по 3 элемента
numgroup1 = []                                                 # Происходит разбиение отфильтрованных чисел на группы по 3 элемента.
numgroup2 = []                                                 # В цикле for каждое число добавляется в numgroup2, и если размер numgroup2 достигает 3
for number in filternum:                                       # группа добавляется в numgroup1, а numgroup2 сбрасывается.
    numgroup2.append(number)                                   #Если в последней группе numgroup2 осталось менее 3 элементов
    if len(numgroup2) == 3:                                    # они добавляются в отдельную группу.
        numgroup1.append(numgroup2)
        numgroup2 = []

# Если в последней группе осталось менее 3 элементов, добавляем их в отдельную группу
if numgroup2:
    numgroup1.append(numgroup2)

# Выводим результат в виде таблицы
for group in numgroup1:                                         #Цикл for перебирает все группы group из numgroup1 и выводит их на печать.
    print(*group)
print('###########################################')
#ЗАдание3

car_brands = ["Toyota", "Honda", "BMW"]
car_numbers = ["ABC123", "DEF456", "GHI789"]
car_owners = [["John"], ["Alice", "Bob"], ["Emily", "David","Sin"]]

# Объединяем списки в зип
car_data = zip(car_brands, car_owners, car_numbers)
for data in car_data:
    print(*data)

# Добавление новой машины
new_brand = input("Введите марку новой машины: ")
new_number = input("Введите номер новой машины: ")
new_owners = input("Введите список владельцев новой машины (через запятую): ").split(",")

car_brands.append(new_brand)
car_numbers.append(new_number)
car_owners.append(new_owners)

car_data = zip(car_brands, car_owners, car_numbers)
for data in car_data:
    print(*data)
# Удаление машины по номеру
remove_number = input("Введите номер машины для удаления: ")
if remove_number in car_numbers:
    index = car_numbers.index(remove_number)
    car_brands.pop(index)
    car_numbers.pop(index)
    car_owners.pop(index)
    print("Машина успешно удалена.")
else:
    print("Машина с таким номером не найдена.")

car_data = zip(car_brands, car_owners, car_numbers)
for data in car_data:
    print(*data)
# Поиск машины по номеру
search_number = input("Введите номер машины для поиска: ")
if search_number in car_numbers:
    index = car_numbers.index(search_number)
    print("Марка авто:", car_brands[index])
    print("Список владельцев:", car_owners[index])
else:
    print("Машина с таким номером не найдена.")


import sys
sys.path.append('../src/')

import torch
from optimizer import minimizeLoss
from torch.autograd import Variable

def computeCost(parameters, input):
    weight, bias = parameters
    return torch.add(torch.matmul(input, weight), bias)

input = torch.randn(1, 10)
weight = Variable(torch.randn(10, 1), requires_grad = True)
bias = Variable(torch.randn(1, 1), requires_grad = True)

parameters = [weight, bias]
optimizer = torch.optim.Adam(parameters, lr = 1)
lossArgs = (parameters, input)

print(computeCost(parameters, input))
minimizeLoss(100, optimizer, computeCost, lossArgs)
print(computeCost(parameters, input))


# 100/100

movie = input()

standardTickets = 0
studentTickets = 0
kidTickets = 0

movieTickets = 0

totalTickets = 0

while movie != "Finish":
    movieCapacity = int(input())
    for i in range(movieCapacity):
        ticket = input()
        if ticket == "standard":
            standardTickets += 1
        elif ticket == "student":
            studentTickets += 1
        elif ticket == "kid":
            kidTickets += 1
        elif ticket == "End":
            break

        movieTickets += 1
        totalTickets += 1

    movieCapacity = movieTickets / movieCapacity * 100
    print(f"{movie} - {movieCapacity:.2f}% full.")
    movieTickets = 0
    movie = input()

standardPercentage = standardTickets / totalTickets * 100
studentPercentage = studentTickets / totalTickets * 100
kidPercentage = kidTickets / totalTickets * 100

print(f"Total tickets: {totalTickets}")
print(f"{studentPercentage:.2f}% student tickets.")
print(f"{standardPercentage:.2f}% standard tickets.")
print(f"{kidPercentage:.2f}% kids tickets.")


from tokenizer import *
from logger import *
import sys

operator_precedence = { # TODO
    "=": 0,
    "+=": 0,
    "-=": 0,
    "*=": 0,
    "**=": 0,
    "/=": 0,
    "%=": 0,
    "==": 1,
    "!=": 1,
    ">": 1,
    "<": 1,
    ">=": 1,
    "<=": 1,
    "+": 3,
    "-": 3,
    "%": 4,
    "*": 5,
    "/": 5,
    "**": 6,
    "!": 7,
    "#": 8,
    "++": 9,
    "--": 9,
    "@": 10,
    "@@": 11,
    ".": 12
}

right_associative = [0, 10]

prefix = ["!", "++", "--", "+", "-", "*", "#"]
infix = ["=", "+=", "-=", "*=", "**=", "/=", "%=", "==", "!=", ">", "<", ">=", "<=", "@", "+", "-", "%", "*", "/", "**", "."]
postfix = ["++", "--", "!", "@@"]

id = -1
def getID():
    global id
    id += 1
    return id

def keepBranch(node, current):
    nodeprec = operator_precedence[node]
    currprec = operator_precedence[current]
    return nodeprec > currprec or nodeprec == currprec and nodeprec not in right_associative

class ParseTree():
    def __init__(self, type = None, content = None, root = None, children = None, flag = False, id = -1):
        self.type = type
        self.content = content
        self.root = root
        self.children = children or []
        self.target = self
        self.id = (id + 1 or getID() + 1) - 1
        self.flag = flag
    def add_child(self, child):
        self.children.append(child)
        child.root = self
    def remove_child(self, child):
        try:
            self.children.remove(child)
            child.root = None
        except:
            pass # Safe if the parameter wasn't a child of this tree anyway
    def pushOperator(self, operator):
        child = ParseTree(type = "operator", content = operator)
        if self.target.type == "operator":
            if operator in postfix and operator not in prefix + infix:
                if self.target.children:
                    self.target.children[-1].add_child(child)
                elif self.target.content in postfix:
                    self.target.add_child(child)
                    self.target = child
                else:
                    raise RuntimeError("Postfix operator directly after non-postfix operator")
            if not self.target.flag and operator in infix and self.target.content in infix and len(self.target.children) == 2:
                if keepBranch(self.target.content, operator):
                    root = self.target.root
                    root.remove_child(self.target)
                    child.add_child(self.target)
                    root.add_child(child)
                    self.target = child
                else:
                    branch = self.target.children[-1]
                    self.target.remove_child(branch)
                    child.add_child(branch)
                    self.target.add_child(child)
                    self.target = child
            elif self.target.content in postfix and operator not in postfix:
                self.target = self.target.root
                self.pushOperator(operator)
            elif not self.target.flag and operator in prefix:
                self.target.add_child(child)
                self.target = child
            elif self.target.flag and operator in infix:
                root = self.target.root
                root.remove_child(self.target)
                child.add_child(self.target)
                root.add_child(child)
                self.target = child
            else:
                raise RuntimeError("No item preceding non-prefix operator")
        elif self.target == self:
            if operator in prefix:
                self.target.add_child(child)
                self.target = child
            else:
                raise RuntimeError("No item preceding non-prefix operator")
        elif self.target.type == "item":
            if operator in postfix:
                self.target.add_child(child)
                self.target = child
            elif (self.target.content in ["list", "bracket"] or self.target.root == self) and (operator in postfix and not self.target.root.flag or operator in prefix and not self.target.flag and self.target.root != self):
                self.target.add_child(child)
                self.target = child
            elif self.target.root.type == "operator":
                if keepBranch(self.target.root.content, operator):
                    root = self.target.root
                    rootroot = self.target.root.root
                    rootroot.remove_child(root)
                    child.add_child(root)
                    rootroot.add_child(child)
                    self.target = child
                else:
                    root = self.target.root
                    root.remove_child(self.target)
                    child.add_child(self.target)
                    root.add_child(child)
                    self.target = child
            else:
                root = self.target.root
                root.remove_child(self.target)
                child.add_child(self.target)
                root.add_child(child)
                self.target = child
        elif self.target.type == "return":
            self.target.add_child(child)
            self.target = child
        elif self.target.type in ["if", "elif", "while", "foreach"]:
            self.target.add_child(child)
    def pushItem(self, item):
        child = ParseTree(type = "item", content = item)
        if item.type == "comma":
            if self.target.type == "item":
                if self.target.root and self.target.root.content == "bracket":
                    self.target.root.flag = True
                    self.target = self.target.root
                else:
                    root = self.target.root
                    root.remove_child(self.target)
                    child = ParseTree(type = "bracket", content = None, root = root, children = [self.target], flag = True, id = -1)
                    root.add_child(child)
                    self.target = child
        elif self.target == self:
            self.target.add_child(child)
            self.target = child
        elif self.target.type == "item":
            self.target = self
            self.pushItem(item)
        elif self.target.type == "operator":
            if self.target.content in prefix and not self.target.children:
                self.target.add_child(child)
                self.target.flag = True
                self.target = self.target.root
            elif self.target.content in infix and len(self.target.children) < 2:
                self.target.add_child(child)
                self.target = child
            elif self.target.content in postfix:
                raise RuntimeError("Cannot push item onto a postfix-only operator")
            else:
                self.target = self.target.root
                self.pushItem(item)
        elif self.target.type in ["return", "if", "elif", "while", "foreach" "import", "include", "pyimport", "pyinclude", "error"]:
            self.target.add_child(child)
            self.target = child
        elif self.target.type == "bracket" and self.target.flag:
            self.target.flag = False
            self.target.add_child(child)
    def getHead(self):
        return self.root.getHead() if self.root else self
    def remove_all(self):
        for child in self.children:
            if child != self:
                child.remove_all()
            self.remove_child(child)
    def __str__(self, target = None, recurse = 0):
        target = target or self.target
        repr = "%s:[Node content = %s type = %s]:%d>%s" % (self.type, str(self.content), type(self.content), self.id, str(self.root.id if self.root else "x")) if self.type else "HEAD:%d%s" % (self.id, "$" if self == target else "")
        # repr = (str(self.content.content) if type(self.content) == type(Token(None, None)) else str(self.content)) if self.type else "HEAD"
        for child in self.children:
            try:
                if recurse > 10:
                    raise RuntimeError("force_overflow")
                lines = child.__str__(target, recurse + 1).split("\n")
                repr += "\n\\----%s%s%s" % (lines[0], "$" if child == target else "", "!" if child.flag else "")
                for line in lines[1:]:
                    repr += "\n|    %s" % line
            except Exception as e:
                repr += "\n    | " + str(e)
        return repr # if recurse or repr.count("$") == 1 else "This is going to get ugly when you try to evaluate it..."
    def __repr(self):
        return str(self)

class Parser():
    def __init__(self, tokens = None):
        self.tokens = tokens or []
        log("Created parser with tokens %s" % str(self.tokens))
    def fill(self):
        tree = ParseTree(None, None, None, [], False, -1)
        index = 0
        balanceable = {
            "listhead": ("listtail", "list"),
            "brakhead": ("braktail", "bracket"),
            "funchead": ("bloctail", "function"),
            "lambdahead": ("bloctail", "lambda"),
            "thenhead": ("bloctail", "then", ["if", "elif"]),
            "elsehead": ("bloctail", "else", ["if", "elif"]),
            "dohead": ("bloctail", "do", ["while", "foreach"])
        }
        ends = ["listtail", "braktail", "bloctail"]
        typed = ["else"]
        while index < len(self.tokens):
            token = self.tokens[index]
            log("------------\nPushing %s onto target:\n%s\n\n" % (str(token), str(tree.target)))
            if token.type == "operator":
                tree.pushOperator(token.content)
            elif token.type in balanceable:
                index += 1
                inner_tokens_list = [[]]
                brackets = 1
                while brackets:
                    if self.tokens[index].type in balanceable:
                        brackets += 1
                        if brackets:
                            inner_tokens_list[-1].append(self.tokens[index])
                    elif self.tokens[index].type in ends:
                        brackets -= 1
                        if brackets:
                            inner_tokens_list[-1].append(self.tokens[index])
                        elif self.tokens[index].type != balanceable[token.type][0]:
                            raise RuntimeError("Mismatched brackets")
                    elif self.tokens[index].type == "comma" and brackets == 1:
                        inner_tokens_list.append([])
                    else:
                        inner_tokens_list[-1].append(self.tokens[index])
                    index += 1
                log("Bracket closed with tokens: %s" % str(inner_tokens_list))
                if balanceable[token.type][2:]:
                    while tree.target.type not in balanceable[token.type][2]:
                        tree.target = tree.target.root
                if balanceable[token.type][1] in typed:
                    child = ParseTree(type = balanceable[token.type][1], content = None, root = None, children = [], flag = False, id = -1)
                else:
                    child = ParseTree(type = "item", content = balanceable[token.type][1], root = None, children = [], flag = False, id = -1)
                tree.target.add_child(child)
                tree.target = child
                for inner_tokens in inner_tokens_list:
                    if inner_tokens:
                        parser = Parser(inner_tokens)
                        subtree = parser.fill()
                        if token.type in ["funchead", "lambdahead", "thenhead", "elsehead", "dohead"]:
                            for subchild in subtree.children:
                                child.add_child(subchild)
                        else:
                            if len(subtree.children) > 1:
                                raise RuntimeError("Subtree has %d children but should have at most 1" % len(subtree.children))
                            if subtree.children:
                                child.add_child(subtree.children[0])
                    else:
                        child.add_child(ParseTree(type = "null"))
                index -= 1
            elif self.tokens[index].type == "statement":
                tree.target = tree
            elif self.tokens[index].type == "keyword":
                requiredParent = {
                    "elif": ["if", "elif"],
                    "do": ["while", "foreach"]
                }
                if self.tokens[index].content in ["if", "return", "elif", "while", "foreach", "import", "include", "pyimport", "pyinclude", "error"]:
                    if self.tokens[index].content in requiredParent:
                        while tree.target.type not in requiredParent[self.tokens[index].content] and requiredParent[self.tokens[index].content]:
                            tree.target = tree.target.root
                    child = ParseTree(self.tokens[index].content)
                    tree.target.add_child(child)
                    tree.target = child
            else:
                tree.pushItem(token)
            log(tree)
            index += 1
        return tree
    def addToken(self, token):
        self.tokens.append(token)

if __name__ == "__main__":
    while True:
        try:
            code = sys.argv[1] if sys.argv[1:] else input("Enter code to be parsed >>> ")

            if sys.argv[1:]: sys.argv = sys.argv[1:]

            tokenizer = Tokenizer(code)

            parser = Parser([])

            while tokenizer.hasNext():
                next = tokenizer.next()
                if next:
                    parser.addToken(next)

            tree = parser.fill()

            print(tree)

            tree.remove_all()
        except Exception as e:
            err("ERROR: %s" % str(e))
            raise


from app import content

CMS_DICT = content()
radio = CMS_DICT['radio']
tv = CMS_DICT['tv']

def get_list(station):
    counter = 0
    for item in station:
        counter += 1
        print(counter, item[0])

if __name__ == '__main__':
    get_list(tv)
    print()
    get_list(radio)


import logging

import matplotlib.pyplot as plt

from kuibit import argparse_helper as kah
from kuibit.simdir import SimDir
from kuibit.visualize_matplotlib import (
    get_figname,
    save_from_dir_filename_ext,
    set_axis_limits_from_args,
    setup_matplotlib,
)

if __name__ == "__main__":
    desc = f"""\
{kah.get_program_name()} plots the times intervals at which the different given apparent horizons were found. """

    parser = kah.init_argparse(desc)
    kah.add_figure_to_parser(parser, add_limits=True)

    parser.add_argument(
        "-a",
        "--horizons",
        type=int,
        required=True,
        help="Apparent horizons to plot",
        nargs="+",
    )

    args = kah.get_args(parser)
    setup_matplotlib(rc_par_file=args.mpl_rc_file)

    # Parse arguments

    logger = logging.getLogger(__name__)

    if args.verbose:
        logging.basicConfig(format="%(asctime)s - %(message)s")
        logger.setLevel(logging.DEBUG)

    horizons = "_".join([str(h) for h in args.horizons])
    figname = get_figname(args, default=f"ah_{horizons}_found")
    logger.debug(f"Figname: {figname}")

    with SimDir(
        args.datadir,
        ignore_symlinks=args.ignore_symlinks,
        pickle_file=args.pickle_file,
    ) as sim:
        logger.debug("Prepared SimDir")

        sim_hor = sim.horizons

        logger.debug(
            f"Apparent horizons available: {sim_hor.available_apparent_horizons}"
        )

        for ah in args.horizons:
            if ah in sim_hor.available_apparent_horizons:
                logger.debug(f"Reading horizon {ah}")
                # We can use any index for the qlm index, it will be thrown away
                current_horizon = sim_hor.get_apparent_horizon(ah)
                time_found = current_horizon.ah.cctk_iteration.t
                # We prepare an array with the same length of time_found and with
                # constant value of ah
                ah_num = [ah] * len(time_found)
                plt.scatter(time_found, ah_num, marker="o", s=0.1)

        # Plot
        logger.debug("Plotting")
        plt.ylabel("Apparent horizon")
        plt.xlabel("Time")

        # Fix ticks
        plt.gca().tick_params(axis="y", which="minor", left=False)
        plt.ylim(min(args.horizons) - 1, max(args.horizons) + 1)
        plt.yticks(args.horizons)

        set_axis_limits_from_args(args)
        logger.debug("Plotted")

        logger.debug("Saving")
        save_from_dir_filename_ext(
            args.outdir,
            figname,
            args.fig_extension,
            tikz_clean_figure=args.tikz_clean_figure,
        )
        logger.debug("DONE")


import socket

# Server IP, PORT
PORT = 8084
IP = "172.20.10.6"

# Before connecting to the server, ask the user for the string
message = "ACTG" + "\n" + "reverse"

# We create sockets for communicating with the server, and then we define how we want it to work: as a file.
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

# Connect to the server
s.connect((IP, PORT))

# Send the request message to the server
s.send(str.encode(message))

# We receive the information from the server, and then we print it
response = s.recv(2048).decode("utf-8")

# Print server response
print("Response: {}".format(response))


import numpy as np
import pygame

from common import window, colors, DEFAULT_FONT, quitProgram
from drawing import drawTarget, drawTracker, writeText
from fuzzy_tracker import FuzzyTracker

pygame.display.set_caption("Fuzzy Logic Target Tracking by Jelena Ristić and Mihailo Grbić")

window_dim = np.array(window.get_size())
freeze = False

tracker = FuzzyTracker(position=window_dim / 2)

while True:
    window.fill(colors['White'])

    # Event handling
    events = pygame.event.get()
    for event in events:
        if event.type == pygame.QUIT:
            quitProgram()

        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_ESCAPE:
                quitProgram()

        if event.type == pygame.MOUSEBUTTONDOWN:
            freeze = not freeze

    # Make target follow mouse pointer if not frozen
    if not freeze:
        target_pos = np.array(pygame.mouse.get_pos())

    tracker.updateTrackerPos(target_pos)

    drawTarget(target_pos)
    drawTracker(tracker.position)
    writeText()

    # Draw everything on the screen
    pygame.display.update()


import random

with open("/Applications/Adobe InDesign CC 2018/Resources/AMT/ID/AMT/application.xml", "r") as f:

    allLines = f.readlines()

trialLine = allLines[21]
keyLength = 0
oldKey = ""

for i in trialLine:

    try:
        currentLetter = int(i)
        keyLength += 1
        oldKey += i
    except:
        pass

#keyLength = 24
#print(oldKey)

with open("/Users/ahish/Desktop/Ahish/Scripts/TrialRestarter/usedKeys", "r") as f:
    allUsedString = f.readlines()

allUsedList = []
for i in allUsedString:
    allUsedList += [int(i.rstrip("\n"))]
'''
new = random.randint(10**(keyLength-1), 10**(keyLength)-1)
print(new)
'''
new = oldKey
while (new in allUsedList ) or (new == oldKey):
    lastDigit = random.randint(0,9)
    while (lastDigit == int(oldKey[-1])):
        lastDigit = random.randint(0,9)
    secondLast = random.randint(0,9)
    while (secondLast == int(oldKey[-1])):
        secondLast = random.randint(0,9)

    new = list(str(new))
    new[-1] = str(lastDigit)
    new[-2] = str(secondLast)
#    print(new)
    temp = ""
    for i in new:
        temp += i
    new = int(temp)
#print(new)

while new in allUsedList:

    new = random.randint(10**(keyLength-1), 10**(keyLength)-1)

newTrialLine = "        <Data key=\"TrialSerialNumber\">"+str(new)+"</Data>\n"

allLines[21] = newTrialLine

with open("/Applications/Adobe InDesign CC 2018/Resources/AMT/ID/AMT/application.xml", "w") as f:

    for i in allLines:
        f.write(i)

with open("/Users/ahish/Desktop/Ahish/Scripts/TrialRestarter/usedKeys", "w") as f:

    for i in allUsedString:
        f.write(i)

    f.write(str(new)+"\n")


#!/usr/bin/env python

import matplotlib.pyplot as plt
import json

def main():
    filename = "3D.txt"

    values = []
    with open(filename, "r") as f:
        line = f.readline()
        while line:
            rho, theta, length = line.split(',')
            values.append((float(rho), float(theta), float(length)))
            line = f.readline()

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    for rho, theta, length in values:
        ax.scatter(rho, theta, length)

    ax.set_xlabel("rho")
    ax.set_ylabel("theta")
    ax.set_zlabel("length")

    plt.show()

if __name__ == "__main__":
    main()


import numpy as np

from common.python.utils import log_utils
from kernel.utils import base_operator

LOGGER = log_utils.get_logger()

def load_data(data_instance):
    X = []
    Y = []
    for iter_key, instant in data_instance:
        weighted_feature = instant.weight * instant.features
        X.append(weighted_feature)
        if instant.label == 1:
            Y.append([1])
        else:
            Y.append([-1])
    X = np.array(X)
    Y = np.array(Y)
    return X, Y

class LogisticGradient(object):

    @staticmethod
    def compute_loss(values, coef, intercept):
        X, Y = load_data(values)
        ywx = np.multiply(-Y.transpose(), X.dot(coef) + intercept).astype(float)
        tot_loss = np.log(1 + np.exp(ywx)).sum()
        return tot_loss

    @staticmethod
    def compute_gradient(values, coef, intercept, fit_intercept):
        X, Y = load_data(values)
        batch_size = len(X)
        if batch_size == 0:
            LOGGER.warning("This partition got 0 data")
            return None
        # [1/(1+e^(-y(wx+b)))-1]*yx
        ywx = -np.multiply(Y.transpose(), X.dot(coef) + intercept).astype(float)
        d = (1.0 / (1 + np.exp(ywx)) - 1).transpose() * Y
        grad_batch = d * X
        if fit_intercept:
            grad_batch = np.c_[grad_batch, d]
        grad = sum(grad_batch)
        return grad

class TaylorLogisticGradient(object):

    @staticmethod
    def compute_gradient(values, coef, intercept, fit_intercept):
        LOGGER.debug("Get in compute_gradient")
        X, Y = load_data(values)
        batch_size = len(X)
        if batch_size == 0:
            return None

        one_d_y = Y.reshape([-1, ])
        d = (0.25 * np.array(base_operator.dot(X, coef) + intercept).transpose() + 0.5 * one_d_y * -1)

        grad_batch = X.transpose() * d
        grad_batch = grad_batch.transpose()
        if fit_intercept:
            grad_batch = np.c_[grad_batch, d]
        grad = sum(grad_batch)
        LOGGER.debug("Finish compute_gradient")
        return grad


# -*- Mode:Python; indent-tabs-mode:nil; tab-width:4 -*-
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import logging

import fixtures
from testtools import TestCase
from testtools.matchers import Equals

from snapcraft_legacy.internal import lifecycle, project_loader, states, steps
from snapcraft_legacy.project import Project
from snapcraft_legacy.storeapi.errors import SnapNotFoundError
from snapcraft_legacy.storeapi.info import SnapInfo
from tests.legacy import fixture_setup

class TestGlobalState(TestCase):
    def setUp(self):
        super().setUp()

        self.useFixture(fixtures.FakeLogger(level=logging.ERROR))

        temp_cwd = fixture_setup.TempCWD()
        self.useFixture(temp_cwd)

        snapcraft_yaml = fixture_setup.SnapcraftYaml(
            temp_cwd.path, base="core20", parts={"test-part": {"plugin": "nil"}}
        )
        self.useFixture(snapcraft_yaml)

        project = Project(
            snapcraft_yaml_file_path=snapcraft_yaml.snapcraft_yaml_file_path
        )

        self.global_state_filepath = project._get_global_state_file_path()

        self.project_config = project_loader.load_config(project)

        self.useFixture(
            fixtures.MockPatchObject(
                self.project_config, "get_build_snaps", return_value={"core20"}
            )
        )

        self.useFixture(
            fixtures.MockPatch(
                "snapcraft_legacy.internal.lifecycle._runner._Executor.run"
            )
        )

        self.useFixture(
            fixtures.MockPatch("snapcraft_legacy.internal.repo.snaps.install_snaps")
        )

        # Avoid unnecessary calls to info.
        channel_map = []
        for arch in ("amd64", "i386", "s390x", "arm64", "armhf", "ppc64el"):
            channel_map.append(
                {
                    "channel": {
                        "architecture": arch,
                        "name": "stable",
                        "released-at": "2019-10-15T13:54:06.800280+00:00",
                        "risk": "stable",
                        "track": "latest",
                    },
                    "confinement": "strict",
                    "download": {
                        "deltas": [],
                        "sha3-384": "64d232d6bfa65be14d7f8d84e952d4e372e12021e2c3dbaf70cf2af5e78bf51c4baf9c9107dd6db815064636b781bda6",
                        "size": 57151488,
                        "url": "https://api.snapcraft.io/api/v1/snaps/download/CSO04Jhav2yK0uz97cr0ipQRyqg0qQL6_1223.snap",
                    },
                    "revision": 1223,
                }
            )
        info = {
            "channel-map": channel_map,
            "default-track": None,
            "name": "core20",
            "snap": {
                "name": "core20",
                "publisher": {
                    "display-name": "Canonical",
                    "id": "canonical",
                    "username": "canonical",
                    "validation": "verified",
                },
                "snap-id": "CSO04Jhav2yK0uz97cr0ipQRyqg0qQL6",
            },
            "snap-id": "CSO04Jhav2yK0uz97cr0ipQRyqg0qQL6",
        }
        self.fake_storeapi_get_info = fixtures.MockPatch(
            "snapcraft_legacy.storeapi._snap_api.SnapAPI.get_info",
            return_value=SnapInfo(info),
        )
        self.useFixture(self.fake_storeapi_get_info)

    def test_stable_grade_for_stable_base(self):
        lifecycle.execute(steps.PULL, self.project_config)

        global_state = states.GlobalState.load(filepath=self.global_state_filepath)
        self.assertThat(global_state.get_required_grade(), Equals("stable"))
        self.fake_storeapi_get_info.mock.assert_called_once_with("core20")

    def test_stable_grade_for_non_stable_base(self):
        self.fake_storeapi_get_info.mock.side_effect = SnapNotFoundError(
            snap_name="core20"
        )
        lifecycle.execute(steps.PULL, self.project_config)

        global_state = states.GlobalState.load(filepath=self.global_state_filepath)
        self.assertThat(global_state.get_required_grade(), Equals("devel"))
        self.fake_storeapi_get_info.mock.assert_called_once_with("core20")

    def test_grade_not_queried_for_if_already_set(self):
        # Set the grade
        global_state = states.GlobalState()
        global_state.set_required_grade("devel")
        global_state.save(filepath=self.global_state_filepath)

        lifecycle.execute(steps.PULL, self.project_config)

        self.fake_storeapi_get_info.mock.assert_not_called()