# -*- coding: utf-8 -*-

import os
from os import fstat
import matplotlib.pyplot as plt
from matplotlib import rcParams
from matplotlib.cm import ScalarMappable

import sys
import pandas as pd
import numpy as np

plt.style.use('seaborn-white')

rcParams['figure.dpi'] = 300
rcParams['image.cmap'] = 'rainbow'
rcParams['ytick.major.size'] = 3
rcParams['xtick.major.size'] = 3
rcParams['font.size'] = 6
rcParams['lines.linewidth'] = 0.75

#### Parsing functions

def read_FStatus(ip_file,eos):

    """
    Function to parse FStatus file.
    It creates a pandas dataframe
    """

    print('Processing {:s} file'.format(ip_file))

    with open(ip_file) as fs:
        next(fs)
        fs_labels = next(fs)

    fs_labels = fs_labels.split('=')[1]

    if eos == 'ECO2N':
        fs_labels = fs_labels.split()
        fs_header = fs_labels[2:]

    else:
        fs_labels = fs_labels.split(",")
        fs_labels = [item.split('(')[0] for item in fs_labels]
        fs_header = fs_labels[4:]

    fstatus = pd.read_csv(ip_file, skiprows=3, names=fs_labels)
    fstatus = fstatus.apply(pd.to_numeric, errors='coerce')
    fstatus = fstatus.fillna(0)
    fstatus['Pres']/= 1e5

    return fs_header, fstatus

def read_FFlow(ip_file, eos):

    """
    Function to parse FFlow file.
    It creates a pandas dataframe
    """

    print('Processing {:s} file'.format(ip_file))

    with open(ip_file) as ff:
        if eos != 'ECO2N':
            next(ff)
        ff_labels = next(ff)

    if eos == 'ECO2N':
        ff_labels = ff_labels.split()
        r_skip = 1

        ff_header = ff_labels[2:]

    else:
        ff_labels = ff_labels.split('=')[1]
        ff_labels = ff_labels.split(",")
        ff_labels = [item.split('(')[0] for item in ff_labels]
        ff_labels.append('V_mix')
        # add dummy header to column created at the end
        ff_labels.append('trail')
        r_skip = 3

        ff_header = ff_labels[4:]

    fflow = pd.read_csv(ip_file, skiprows=r_skip, names=ff_labels)
    fflow = fflow.apply(pd.to_numeric, errors='coerce')
    fflow = fflow.fillna(0)

    if eos != 'ECO2N':
        fflow = fflow.drop('trail', axis=1)
        ff_header.pop()

    return ff_header, fflow

def read_COFT(ip_file, eos):

    """
    Function to parse coft file.
    It creates a pandas dataframe
    """

    print('Processing {:s} file'.format(ip_file))

    coft = pd.read_csv(ip_file, header=None, index_col=0, low_memory=False)

    coft = coft.apply(pd.to_numeric, errors='coerce')

    #Drop empty columns
    coft = coft.dropna(axis=1, how='all')

    #Replace remaining NaN values for zero
    coft = coft.fillna(0)

    if eos == 'ECO2N':
        coft_var = ['FLO(GAS)', 'FLO(aq.)', 'VEL(GAS)', 'VEL(LIQ.)', 'FLO(NaCl)' , 'FLO(CO2)' ,'FLOH']
        c_idx_df = coft.loc[:,2::8]
    else:
        coft_var = ['FLO(aq.)', 'FLO(LIQ.)', 'FLO(GAS)', 'FLO(CO2)']
        c_idx_df = coft.loc[:,2::5]

    #Retrieve connection indexes and delete columns

    c_idx = c_idx_df.drop_duplicates().values.flatten()
    coft = coft.drop(columns=c_idx_df.columns)

    print('{:4d} connections reported in coft.'.format(len(c_idx)))

    #Define multiIndex to rename columns
    coft_col = pd.MultiIndex.from_product([c_idx,coft_var])
    coft_col = coft_col.insert(loc=0, item='time')

    coft.columns = coft_col

    return coft_var, c_idx, coft

def read_FOFT(ip_file, eos):

    """
    Function to parse foft file.
    It creates a pandas dataframe
    """

    print('Processing {:s} file'.format(ip_file))

    foft = pd.read_csv(ip_file, header=None, index_col=0, low_memory=False)

    foft = foft.apply(pd.to_numeric, errors='coerce')

    #Drop empty columns
    foft = foft.dropna(axis=1, how='all')

    #Replace remaining NaN values for zero
    foft = foft.fillna(0)

    #Define headers
    if eos == 'ECO2N':
        foft_var = ['Pres', 'Sg', 'XNACL', 'XCO2liq', 'T']
    else:
        foft_var = ['Pres', 'T', 'S_liquid', 'S_gas', 'XCO2liq']

    #Retrieve element indexes and delete columns
    e_idx_df = foft.loc[:,2::6]

    e_idx = e_idx_df.drop_duplicates().values.flatten()

    print('{:4d} grid elements reported in foft.'.format(len(e_idx)))

    foft = foft.drop(columns=e_idx_df.columns)

    #Define multiIndex to rename columns
    foft_col = pd.MultiIndex.from_product([e_idx,foft_var])
    foft_col = foft_col.insert(loc=0, item='time')

    foft.columns = foft_col

    foft.iloc[:,1::5] = foft.iloc[:,1::5]/1e5

    return foft_var, e_idx, foft

def read_ipMESH(fname):

    """
    Parse TOUGH input file and extract grid parameters
    from ELEME and CONNE sections into a dataframe
    """

    mesh = dict()

    with open(fname, 'r') as ip_file:
        f = ip_file.read()

        if 'ELEME' in f:
            read_MESH_file = False
        else:
            read_MESH_file = True

    if read_MESH_file:
        print("Read ELEME and CONNE from MESH file")
        with open(r'MESH') as mesh_file:
            f = mesh_file.read()

    ELEME_i = f.find('ELEME')
    ELEME_f = ELEME_i + f[ELEME_i:].find('CONNE')
    ELEME_raw = f[ELEME_i:ELEME_f].split('\n')[1:]

    ELEME = pd.DataFrame(data = ELEME_raw, index = 1+np.arange(len(ELEME_raw)))
    ELEME = ELEME.drop(index=ELEME[ELEME[0].str.strip().str.len()==0].index)

    ELEME['ElName'] = ELEME[0].str[:5]
    ELEME['NSEQ'] = ELEME[0].str[5:10]
    ELEME['NADD'] = ELEME[0].str[10:15]
    ELEME['MAT'] = ELEME[0].str[15:20]
    ELEME['VOLX'] = pd.to_numeric(ELEME[0].str[20:30], 'coerce')
    ELEME['AHTX'] = pd.to_numeric(ELEME[0].str[30:40], 'coerce')
    ELEME['PMX'] = pd.to_numeric(ELEME[0].str[40:50], 'coerce')
    ELEME['X'] = pd.to_numeric(ELEME[0].str[50:60], 'coerce')
    ELEME['Y'] = pd.to_numeric(ELEME[0].str[60:70], 'coerce')
    ELEME['Z'] = pd.to_numeric(ELEME[0].str[70:80], 'coerce')

    #Remove zero volume elements set for Dirichlet condition
    ELEME.loc[ELEME.VOLX.isna(),'VOLX'] = 0

    #Reset index to preserve correct index values
    ELEME  = ELEME.loc[~(ELEME.VOLX==0)].copy().reset_index(drop=True)
    ELEME.index = ELEME.index+1

    ELEME = ELEME.drop(columns=0)

    CONNE_i = f.find('CONNE')

    if read_MESH_file:
        CONNE_f = CONNE_i + f[CONNE_i:].find('+++')
        CONNE_raw = f[CONNE_i:CONNE_f].split('\n')[1:-1]
    else:
        f = f.replace('\n \n', '\n\n')
        CONNE_f = CONNE_i + f[CONNE_i:].find('\n\n')
        CONNE_raw = f[CONNE_i:CONNE_f].split('\n')[1:]

    CONNE = pd.DataFrame(data = CONNE_raw, index = 1+np.arange(len(CONNE_raw)))
    CONNE['EL1'] = CONNE[0].str[:5]
    CONNE['EL2'] = CONNE[0].str[5:10]
    CONNE['NSEQ'] = CONNE[0].str[10:15]
    CONNE['NAD1'] = CONNE[0].str[15:20]
    CONNE['NAD2'] = CONNE[0].str[20:25]
    CONNE['ISOT'] = pd.to_numeric(CONNE[0].str[25:30], 'coerce')
    CONNE['D1'] = pd.to_numeric(CONNE[0].str[30:40], 'coerce')
    CONNE['D2'] = pd.to_numeric(CONNE[0].str[40:50], 'coerce')
    CONNE['AREAX'] = pd.to_numeric(CONNE[0].str[50:60], 'coerce')
    CONNE['BETAX'] = pd.to_numeric(CONNE[0].str[60:70], 'coerce')
    CONNE['SIGX'] = pd.to_numeric(CONNE[0].str[70:80], 'coerce')

    CONNE = CONNE.drop(columns=0)

    return ELEME, CONNE

#Plotting functions

name_variable = ['Depth', 'Dis', 'cumDepth', 'Dgas', 'D_aqueous' , 'D_liquid', 'D_gas', 'FLO(CO2)', 'FLO(GAS)',
                'FLO(NaCl)', 'FLO(aq.)', 'FLO(LIQ.)',  'q_aqueous', 'q_liquid', 'q_gas' , 'FLOH', 'Fgas',
                'Fliq', 'Pres', 'Sg', 'S_aqueous', 'S_liquid', 'S_gas', 'T', 'Time',
                'Umix', 'VEL(GAS)', 'VEL(LIQ.)', 'VGas',  'V_aqueous', 'V_liquid', 'V_gas', 'V_mix',
                'VLiq', 'XCO2liq', 'XNACL', 'var1', 'var2', 'var3', 'var4', 'WellID']

unit_variable = ['m']*3 +    4*['kg/m$^3$'] + 8*['kg/s'] + ['W'] + 2*['kg/s'] + ['bar'] + ['m$^3$/m$^3$']*4 + ['$\degree$C', 's'] + 9*['m/s'] + 2*['kg/kg'] + 5*['-']
unit_variable_v2 = ['m']*3 + 4*['kg/m3']    + 8*['kg/s'] + ['W'] + 2*['kg/s'] + ['bar'] + ['m3/m3'      ]*4 + ['degC', 's']       + 9*['m/s'] + 2*['kg/kg'] + 5*['-']

unit_lims = {'Sg':(0,1)}

units_dict = dict(zip(name_variable,unit_variable))
units_dict_v2 = dict(zip(name_variable,unit_variable_v2))

perm_dict = {-1:'-x', 1:'+x', -2:'-y', 2:'+y', -3:'-z', 3:'+z'}

def secondary_scale(log_bool, ax):
    """
    Sets a secondary scale with human readable units
    Takes two parameters
    log_bool: Boolean set as True if the scale is logarithmic, False if linear
    ax: The matploltib artist axis

    Returns the input axis with a new secosndary x-axis on top
    """
    if log_bool:
        ax.set_xscale('log')

        sec_x = ax.twiny()

        sec_x.set_xscale('log')
        x_ticks = np.array([1,60,3600,3600*24,3600*24*30, 3600*24*365.25, 3600*24*365.25*100, 3600*24*365.25*1000])
        x_labels = np.array(['1 s', '1 min', '1 h', '1 day', '1 month', '1 yr', '100 yr', '1 kyr'])
        sec_x.set_xticks(x_ticks)
        sec_x.set_xticklabels(x_labels)
        sec_x.set_xlim(ax.get_xlim())
        sec_x.set_xlabel('time')

    else:

        t_min, t_max = ax.get_xlim()

        scales = [60, 3600, 3600*24, 3600*24*365.25, 3600*24*365.25*1000]
        units = ['min', 'hr', 'd', 'yr', 'kyr']

        n_ticks = 5

        fix = np.arange(n_ticks+1)

        sec_x = ax.twiny()
        for s,u in zip(scales, units):
            scaled_tmax = t_max/s

            order = np.log10(scaled_tmax)

            if order>0 and order<=3:
                if order<2:
                        if scaled_tmax/n_ticks<1:

                                step = (5*(1+10*(scaled_tmax/n_ticks))//5)/10
                                x_labels = fix*step

                        elif scaled_tmax<=10:
                                x_labels = np.arange(np.ceil(scaled_tmax))
                        else:
                                step = scaled_tmax//n_ticks
                                x_labels = np.arange(np.ceil(scaled_tmax), step = step)

                else:
                        step = 50*(1+(scaled_tmax//n_ticks)//50)
                        x_labels = fix*step

                x_labels = np.round(x_labels, 1)
                x_ticks = x_labels*s
                sec_unit = u

                sec_x.set_xticks(x_ticks)
                sec_x.set_xticklabels(x_labels)
                sec_x.set_xlim(ax.get_xlim())
                sec_x.set_xlabel('time [{:s}]'.format(sec_unit))

def plot_Ffigure(title,df,df_vars, logscale, eos):
    """
    Function for plotting the FFlow or FStatus files
    Takes three parameters
    title: FStatus or FFlow
    df: the dataframe corresponding to each file
    vars: the variables set for plotting

    Returns and save a figure
    """
    #filter index bars
    index_var = ['WellID',  'Time',  'Dis',  'cumDepth', 'Depth']
    plot_vars = [item for item in df_vars if item not in index_var]

    #Set plot dimensions
    w = 5
    h = (2.5-1)*len(plot_vars)
    rcParams['figure.figsize'] = [w,h]

    #Create axes
    fig, axs = plt.subplots(len(plot_vars),1, sharex=True)

    #Pick right depth vector depending on labels difference in EOS version
    if EOS == 'ECO2N':
        Z_label = 'Depth'
    else:
        Z_label = 'Dis'

    #Define x, y dimensions for creating contour plot
    size_j = df[Z_label].drop_duplicates().shape[0]
    size_i = int(df['Time'].shape[0]/size_j)

    #Define the X and Y axis
    X = df['Time'].values.reshape(size_i, size_j)
    Y = df[Z_label].values.reshape(size_i, size_j)

    for var_idx, var in enumerate(plot_vars):

        #Define limits for saturation variables
        if var in ['Sg', 'S_aqueous', 'S_liquid', 'S_gas']:
            unit_lims[var] = (0,1)

        if len(plot_vars)>1:
            ax = axs[var_idx]
        else:
            ax = axs

        Z = df[var].values.reshape(size_i, size_j)

        try:
            cf = ax.contourf(X, Y, Z, vmin = unit_lims[var][0], vmax = unit_lims[var][1])
            cb = plt.colorbar(ScalarMappable(norm=cf.norm, cmap=cf.cmap), ax=ax,
                                ticks=np.linspace(unit_lims[var][0], unit_lims[var][1], 6))
        except:
            cf = ax.contourf(X, Y, Z)
            cb = plt.colorbar(cf, ax=ax)

        cb.set_label('{:s} [{:s}]'.format(var,units_dict[var]))
        ax.set_ylabel('depth [m]')

        ax.invert_yaxis()

        if var == plot_vars[-1]:
            ax.set_xlabel('time [s]')

        if var == plot_vars[0]:
            secondary_scale(logscale, ax)

    fig.suptitle(title)
    fig.align_ylabels()
    fig.tight_layout(rect=[0,0,1,0.98])
    fig.savefig(r'fig_{:s}.png'.format(title))

def plot_OFT(title, df, items, df_vars, logscale, mesh_eleme, mesh_conne):
    """
    Function for plotting the COFT or FOFT files
    Takes four parameters
    title: COFT or FOFT
    df: the dataframe corresponding to each file
    items: The connection(s) or element(s) set for plottinh
    vars: the variables set for plotting

    Returns and save a figure
    """
    w = 6
    h = (2.8-1)*len(df_vars)

    w = 8.1
    h = 5.85
    rcParams['figure.figsize'] = [w,h]

    mesh_eleme_v2 = mesh_eleme.copy()
    mesh_eleme_v2 = mesh_eleme_v2.set_index(('ElName'))

    fig, axs = plt.subplots(len(df_vars),1, sharex=True)

    for var_idx, var in enumerate(df_vars):

        print(var_idx, var)

        if len(df_vars)>1:
            ax = axs[var_idx]
        else:
            ax = axs

        for item in items:
            # print('plot item {:d} in {:s} plot'.format(item,var))

            if title =='coft':

                el1 = mesh_conne.loc[item, 'EL1']
                el2 = mesh_conne.loc[item, 'EL2']
                k_dir = mesh_conne.loc[item, 'ISOT']

                mat1 = mesh_eleme_v2.loc[el1,'MAT']
                mat2 = mesh_eleme_v2.loc[el2,'MAT']
                item_label =  '{:s}>{:s}({:s} to {:s} in {:s} dir.)'.format(el1,el2,mat1, mat2, perm_dict[k_dir])

            elif title =='foft':
                el = mesh_eleme.loc[item,'ElName']
                mat = mesh_eleme.loc[item,'MAT']
                item_label = '{:<4d}{:s}({:s})'.format(item,el,mat)
            df.plot(x='time', y=(item, var), ax=ax, label=item_label, legend = False)

        if var == df_vars[-1]:
            ax.set_xlabel('time [s]')

        if var == df_vars[0]:
            secondary_scale(logscale, ax)

        ax.set_ylabel('{:s} [{:s}]'.format(var, units_dict[var]))

    if title == 'foft':
        ax.set_yscale('log')

    fig.suptitle(title)
    handles, labels = ax.get_legend_handles_labels()
    fig.legend(handles, labels, loc='center right')
    fig.align_ylabels()
    fig.tight_layout(rect=[0,0,0.75,0.98])
    fig.savefig(r'fig_{:s}.png'.format(title))

#Define variables for plotting based on input arguments

def get_EOS(fname):
    """Retrieves EOS from output file"""
    with open(fname, 'r', encoding='latin1') as of:
        lines = of.readlines()

    for line in lines:
        line = line.strip()
        if line.startswith('EOS'):
            EOS = line.split()[5].strip('*')
            return EOS

def plot_specs(ip_args, plot_bool, files):
    """Defines the files and variables that will be plotted"""
    plot_dict = dict()

    if len(ip_args)>2:

        #Turn off all plotting
        for f in plot_bool:
            plot_bool[f] = False

        #Based on arguments turn on the selected files to plot

        for arg in ip_args[2:]:

            arg = arg.split('i')
            arg_v = arg[0].strip(",").split(",")

            try:
                arg_i = arg[1].strip(",").split(",")
            except:
                arg_i = None

            print('file queried:', arg_v)
            print('items queried:', arg_i)

            #Check if FStatus is being queried
            if any(item.lower().startswith('fstatus') for item in files) and arg_v[0].lower() == r'fstatus':
                plot_bool[r'fstatus'] = True
                if len(arg_v)>1:
                    plot_dict[r'fstatus'] = arg_v[1:]
                else:
                    plot_dict[r'fstatus'] = 'all'

            #Check if FFLow is being queried
            if any(item.lower().startswith('fflow') for item in files) and arg_v[0].lower() == r'fflow':

                plot_bool[r'fflow'] = True

                if len(arg_v)>1:
                    plot_dict[r'fflow'] = arg_v[1:]
                else:
                    plot_dict[r'fflow'] = 'all'

            #Check if coft is being queried
            elif any(item.startswith('COFT') for item in files) and arg_v[0].lower() == r'coft':
                plot_bool[r'coft'] = True
                plot_dict[r'coft'] = dict()
                if len(arg_v)>1:
                    plot_dict[r'coft']['var'] = arg_v[1:]
                else:
                    plot_dict[r'coft']['var'] = 'all'

                if arg_i is None:
                    plot_dict[r'coft']['item'] = 'all'
                else:
                    plot_dict[r'coft']['item'] = arg_i

            #Check if foft is being queried
            elif any(item.startswith('FOFT') for item in files) and arg_v[0].lower() == r'foft':

                plot_bool[r'foft'] = True
                plot_dict[r'foft'] = dict()
                if len(arg_v)>1:
                    plot_dict[r'foft']['var'] = arg_v[1:]
                else:
                    plot_dict[r'foft']['var'] = 'all'

                if arg_i is None:
                    plot_dict[r'foft']['item'] = 'all'
                else:
                    plot_dict[r'foft']['item'] = arg_i

    else:

        for f in plot_bool:

            plot_dict[f] = 'all'

            if f in ['coft', 'foft']:
                plot_dict[f]=dict()
                plot_dict[f]['var'] = 'all'
                plot_dict[f]['item'] = 'all'

    return plot_bool, plot_dict

if __name__ == '__main__':
    args = sys.argv
    print(args)
    logscale = False

    ip_path = args[1]
    ip_dirname = os.path.dirname(ip_path)

    old_path = os.getcwd()

    if len(ip_dirname)>0:
        os.chdir(ip_dirname)

    raw_names = tuple([r'fflow', r'fstatus', r'coft', r'foft'])
    fnames = []
    plot_bool = dict()
    fnames_map=dict()

    for f in os.listdir():
        if f.lower().startswith(raw_names):
            f_size = os.path.getsize(f)
            if f_size>0:
                fnames.append(f)
                flabel = f.lower().split('_')[0]
                fnames_map[flabel] = f
                plot_bool[flabel] = True

        elif f.endswith(tuple(['.in', '.inp'])):
            ip_file = f

        elif f.endswith('out'):
            op_file = f

    EOS = get_EOS(op_file)
    print('Input file: {:s}'.format(ip_file))
    print('EOS version: {:s}'.format(EOS))

    #Define horizontal scale type
    if 'log' in args:
        logscale = True
        args.remove('log')

    #Define how files will be parsed
    parse_dict = dict()

    for file in fnames_map:
        if file.lower().startswith('fflow'):
            parse_dict[file]=read_FFlow
        elif file.lower().startswith('fstatus'):
            parse_dict[file]=read_FStatus
        elif file.startswith('coft'):
            parse_dict[file]=read_COFT
        elif file.startswith('foft'):
            parse_dict[file]=read_FOFT

    #Define which files and variables will be plotted
    plot_bool, plot_dict = plot_specs(args, plot_bool, fnames)
    print(plot_bool)
    print()
    print(plot_dict)

    eleme, conne = read_ipMESH(ip_file)

    #Query and index data

    for ftype in fnames_map:

        file = fnames_map[ftype]

        plot_f = plot_bool[ftype]

        if plot_f:
            print(f'Plotting {file} data\n')

            if ftype in ['fflow', 'fstatus']:
                queried_vars = plot_dict[ftype]

                df_vars, df = parse_dict[ftype](file, EOS)

                selected_var = df_vars

                if queried_vars != 'all':
                    vars_sel = []
                    for var in queried_vars:
                        try:
                            var_idx = (int(var)-1)%len(df_vars)
                            vars_sel.append(df_vars[var_idx])

                        except:
                            for var_name in df_vars:
                                if var.lower() == var_name.lower():
                                    vars_sel.append(var_name)
                    selected_var = vars_sel

                print(f'Plot for "{file}" includes: {" ".join(selected_var)}')

                plot_Ffigure(ftype, df,selected_var, logscale, EOS)

            elif ftype in ['coft', 'foft']:
                print(plot_dict)

                queried_vars = plot_dict[ftype]['var']
                queried_items = plot_dict[ftype]['item']

                # print(queried_vars)

                df_vars, df_items, df = parse_dict[ftype](file, EOS)

                selected_var = df_vars
                selected_items = df_items

                # print(selected_items)

                if queried_vars != 'all':
                    vars_sel = []
                    for var in queried_vars:
                        try:
                            var_idx = (int(var)-1)%len(df_vars)
                            vars_sel.append(df_vars[var_idx])

                        except:
                            for var_name in df_vars:
                                if var.lower() == var_name.lower():
                                    vars_sel.append(var_name)
                    selected_var = vars_sel

                if queried_items != 'all':
                    item_sel = []
                    for item in queried_items:
                        try:
                            item_idx = (int(item)-1)%len(df_items)
                            item_sel.append(df_items[item_idx])

                        except:
                            for item_name in df_items:
                                if item.lower() == item_name.lower():
                                    item_sel.append(item_name)
                    selected_items = item_sel

                print('{:s} plotted variables include: {:s}'.format(file, ' '.join(selected_var)))
                print('{:s} plotted items include: {:s}'.format(file, ' '.join(map(str,selected_items))))
                plot_OFT(ftype, df, selected_items, selected_var, logscale, mesh_eleme=eleme, mesh_conne=conne)

    spreadsheet = ip_file.split(".")[0]+".xlsx"

    print_Excel = True

    xls_output = input('Do you want to store output files as spreadsheet (Y/N). Default is Yes?\t')

    if xls_output.lower().startswith('n'):
        print_Excel = False

    if print_Excel:
        pd_units = pd.Series(units_dict_v2)

        if r'fstatus' in fnames_map.keys():
            #Add column names to FStatus

            fs_var, fs = read_FStatus(fnames_map['fstatus'], EOS)
            fs_row1 = fs.columns.to_list()
            fs_row2 = pd_units[fs_row1].to_list()

            fs_cols = pd.MultiIndex.from_tuples(list(zip( fs_row1, fs_row2)))

            fs.columns = fs_cols

        if r'fflow' in fnames_map.keys():

            #Add column names to FFlow

            ff_var, ff = read_FFlow(fnames_map['fflow'], EOS)
            ff_row1 = ff.columns.to_list()
            ff_row2 = pd_units[ff_row1].to_list()

            ff_cols = pd.MultiIndex.from_tuples(list(zip( ff_row1, ff_row2)))

            ff.columns = ff_cols

        if r'coft' in fnames_map.keys():

            #Add column names to coft
            coft_var, coft_idx, coft = read_COFT(fnames_map['coft'], EOS)

            eleme2 = eleme.copy()
            eleme2 = eleme2.set_index('ElName')

            query_x = conne['ISOT']==1
            query_y = conne['ISOT']==2
            query_z = conne['ISOT']==3

            coft_row1 = coft.columns.get_level_values(0).to_list()
            coft_row2 = (conne.loc[coft_row1[1:],'EL1']+conne.loc[coft_row1[1:],'EL2']).to_list()
            coft_row2 = [''] + coft_row2
            coft_row3 = coft.columns.get_level_values(1).to_list()
            coft_row4 = ['s'] + list(pd_units[coft.columns.get_level_values(1)[1:].to_list()].values)

            coft_cols = pd.MultiIndex.from_tuples(list(zip( coft_row1,
                                                            coft_row2,
                                                            coft_row3,
                                                            coft_row4)))

            coft.columns = coft_cols

        if r'foft' in fnames_map.keys():

            #Add column names to foft
            foft_var, foft_idx, foft = read_FOFT(fnames_map['foft'], EOS)

            foft_row1 = foft.columns.get_level_values(0).to_list()
            foft_row1[0] = 'cell_idx'

            foft_row2 = eleme.loc[foft_row1[1:],'ElName'].to_list()
            foft_row2 = ['cell_name'] + foft_row2

            foft_row2a = ['cell_X [m]'] + eleme.loc[foft_row1[1:], 'X'].to_list()
            foft_row2b = ['cell_Y [m]'] + eleme.loc[foft_row1[1:], 'Y'].to_list()
            foft_row2c = ['cell_Z [m]'] + eleme.loc[foft_row1[1:], 'Z'].to_list()
            foft_row2d = ['cell_mat'] + eleme.loc[foft_row1[1:], 'MAT'].to_list()

            foft_row3 = foft.columns.get_level_values(1).to_list()
            foft_row3[0] = 'time'
            foft_row4 = ['s'] + list(pd_units[foft.columns.get_level_values(1)[1:].to_list()].values)

            foft_cols = pd.MultiIndex.from_tuples(list(zip( foft_row1,
                                                            foft_row2,
                                                            foft_row2a,
                                                            foft_row2b,
                                                            foft_row2c,
                                                            foft_row2d,
                                                            foft_row3,
                                                            foft_row4)))

            foft.columns = foft_cols

        with pd.ExcelWriter(spreadsheet) as writer:
            if r'fstatus' in fnames_map.keys():
                fs.to_excel(writer, sheet_name=r'FStatus')
            if r'fflow' in fnames_map.keys():
                ff.to_excel(writer, sheet_name=r'FFlow')
            if r'coft' in fnames_map.keys():
                coft.to_excel(writer, sheet_name=r'coft')
            if r'foft' in fnames_map.keys():
                foft.to_excel(writer, sheet_name=r'foft')


#!/usr/local/bin/python3
# -*- coding: utf-8 -*-

import os

import numpy as np
import pandas as pd

from sklearn import metrics
from sklearn.model_selection import KFold
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression

import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt

"""
paper figure 6.
"""

def generate_dataset_from_csv(target_name='Ea', fea_names='all'):
    """
    Description:
        Return a dictionary with mechanism types, target(activation energy),
        features(adsorption energies and geometrical descriptors) and their names.
    """
    # from csv file to dataframe
    if os.path.exists(DS_CSV):
        df = pd.read_csv(DS_CSV, index_col=0)
    else:
        raise ValueError('DataSet CSV does not exist.')

    # choose different targets: preferred or specific activation energy (three kinds)
    if target_name == 'Ea':
        df = df
        E = 'mE'
    elif target_name == 'Edelta':
        df = df.loc[df.loc[:,'E_ts']!='np.nan', :]
        df = df.loc[df.loc[:,'E_tsra']!='np.nan', :]
        df.loc[:,'E_ts'] = df.loc[:,'E_ts'].values.astype(np.float64) - \
                df.loc[:,'E_tsra'].values.astype(np.float64)
        E = 'E_ts'
    elif target_name == 'Ets':
        df = df.loc[df.loc[:,'E_ts']!='np.nan', :]
        E = 'E_ts'
    elif target_name == 'Etsra':
        df = df.loc[df.loc[:,'E_tsra']!='np.nan', :]
        E = 'E_tsra'

    # 'name', 'mtype', 'E_ts', 'E_tsra', Ea: min(E_ts, E_tsra)
    nodescriptors = ['name', 'mtype', 'E_ts', 'E_tsra', 'Ea']
    indexs_cols = df.iloc[:, 0:len(nodescriptors)]

    # choose number of descriptors (E_Hab3, E_CH3ab, Geos, ...)
    if fea_names == 'all':
        vals_cols = df.iloc[:,range(len(nodescriptors),len(df.columns))]
    else:
        vals_cols = df.loc[:, fea_names]

    # Get DataSet
    DS = {}

    DS['Etype'] = indexs_cols.loc[:, 'mtype'].values
    DS['target'] = indexs_cols.loc[:,E].values.reshape(-1,1).astype(np.float64)
    DS['fea_names'] = vals_cols.columns.values
    DS['features'] = vals_cols.values.astype(np.float64)

    return DS

def stardard_scaling(y, X):
    """
    Description:
        StandardScaler Data.
    In:
        unnormalized data y and X.
    Out:
        normalized data: y_s and X_s.
        scaler function: s_y and s_X
    """
    # scale y
    scaler4y = StandardScaler().fit(y)
    y_scalered = scaler4y.transform(y)

    # scale X
    scaler4X = StandardScaler()
    X_scalered = scaler4X.fit_transform(X)

    return y_scalered, X_scalered, scaler4y, scaler4X

def randomize_folds(y, splits):
    """"""
    kf = KFold(n_splits=splits, shuffle=True, random_state=0)
    folds_index = []
    for train_index, test_index in kf.split(y):
        folds_index.append([train_index, test_index])

    return folds_index

def linear_regression_with_cv(y, X):
    # generate k-folds
    n = 5
    folds_index = randomize_folds(y, n)
    MSEs, r2s = [], []
    ols = LinearRegression(fit_intercept=True, normalize=False, n_jobs=-1)
    for i in range(n):
        # get train and test set
        fold = folds_index[i]
        train_index, test_index = fold[0], fold[1]
        y_train, X_train = y[train_index], X[train_index]
        y_test, X_test = y[test_index], X[test_index]

        # scale data
        y_s, X_s, s4y, s4X = stardard_scaling(y_train, X_train)

        # train and test
        ols.fit(X_s, y_s)

        y_test_s, X_test_s = s4y.transform(y_test), s4X.transform(X_test)
        y_p = ols.predict(X_test_s)

        MSE = metrics.mean_squared_error(y_test_s, y_p)
        MSEs.append(MSE)

        r2 = metrics.r2_score(y_test_s, y_p)
        r2s.append(r2)

    # mean MSE and standard MSE
    MSEs = np.array(MSEs)

    return MSEs, r2s

def feature_selection(Etype, features):
    """"""
    # prepare data
    DS = generate_dataset_from_csv(Etype, features.keys())
    y, X = DS['target'], DS['features']

    # feature selection
    mean_MSEs, std_MSEs, mean_r2s, std_r2s = [], [], [], []
    feas_number = len(X[0])

    for i in range(feas_number):
        X_slice = []
        for j in range(i+1):
            X_slice.append(X.T[j])
        X_slice = np.array(X_slice).T
        MSEs, r2s = linear_regression_with_cv(y, X_slice)

        mean_MSEs.append(np.mean(MSEs))
        std_MSEs.append(np.var(MSEs))

        mean_r2s.append(np.mean(r2s))
        std_r2s.append(np.var(r2s))

    return mean_MSEs, std_MSEs, mean_r2s, std_r2s

def plot_selected_features(ax, features, feanames):
    """"""
    # figure general setting
    nfeatures = len(features.keys())

    ax.set_title('Descriptors Selected by Lasso', fontsize=16)

    ax.set_yticks(np.arange(nfeatures))
    ax.set_yticklabels(feanames, fontsize=10, fontweight='black')

    ax.set_xlabel('Coefficient', fontsize=16)
    ax.set_xlim(0, 1.0)

    # horizental bar
    coefs = np.array(list(features.values()))
    pos1 = ax.barh(np.arange(nfeatures), np.fabs(coefs), 0.5, \
            color='b', label='Positive', alpha=0.8)

    neg_features = {}
    for i, coef in enumerate(coefs):
        if coef <= 0e-6:
            neg_features[i] = np.fabs(coef)

    neg1 = ax.barh([int(i) for i in neg_features.keys()], \
            list(neg_features.values()), 0.5, \
            color='r', label='Negative', alpha=0.8)

    for i, coef in enumerate(features.values()):
        ax.text(np.fabs(coef), i, '%.2f' %coef, \
                fontsize=12, fontweight='bold')

    ax.legend(handles=[pos1, neg1], loc=1)

def plot_feature_selection(ax, Etype, features):
    """"""
    ax.set_title('Feature Selection by Linear Regression', fontsize=16)

    ax.set_xlabel('First N Features', fontsize=16)

    ax.set_ylabel('Statistics', fontsize=16)
    ax.set_ylim(-0.2, 1.0)

    mean_MSEs, std_MSEs, mean_r2s, std_r2s = feature_selection(Etype, features)

    mr2 = ax.scatter(range(1, len(mean_r2s)+1), mean_r2s, \
            color='g', marker='o', label='Mean $R^2$')
    mmse = ax.scatter(range(1, len(mean_MSEs)+1), mean_MSEs, \
            color='r', marker='o', label='Mean MSE')

    '''
    sr2 = ax.scatter(range(1, len(std_r2s)+1), std_r2s, \
            color='b', marker='x', label='Variance $R^2$')
    smse = ax.scatter(range(1, len(std_MSEs)+1), std_MSEs, \
            color='r', marker='x', label='Variance MSE')
    '''

    ax.legend(handles=[mr2, mmse], loc=5)

def plot_best_regression(ax, Etype, features, first_n):
    """"""
    DS = generate_dataset_from_csv(Etype, list(features.keys())[:first_n])
    y, X = DS['target'], DS['features']
    y_s, X_s, s4y, s4X = stardard_scaling(y, X)

    ols = LinearRegression(fit_intercept=True, normalize=False, n_jobs=-1)
    ols.fit(X_s, y_s)

    y_p = s4y.inverse_transform(ols.predict(X_s))
    r2 = metrics.r2_score(y, y_p)
    mse = metrics.mean_squared_error(y, y_p)

    ax.set_title('First %d Descriptors Linear Regression' %first_n, fontsize=16)

    ax.set_xlabel('True Values / eV', fontsize=16)
    ax.set_xlim(-0.75, 1.75)

    ax.set_ylabel('Predicted Values / eV', fontsize=16)
    ax.set_ylim(-0.75, 1.75)

    ax.scatter(y, y_p)
    ax.plot([-0.5, 1.5], [-0.5, 1.5], 'k--', lw=4)
    ax.text(-0.5, 1.0, '${R^2}$=%0.2f, MSE=%0.2f' %(r2, mse), \
            fontsize=12)

def plot_main():
    """"""
    # figure general setting
    fig, ax = plt.subplots(2, 3, figsize=(16, 9))
    plt.tight_layout(pad=2.0, w_pad=2.0, h_pad=2.0)
    plt.subplots_adjust(left=0.1, bottom=None, right=0.90, top=0.85, \
            wspace=0.3, hspace=0.5)

    plt.suptitle('Feature Selection of Geometrical Descriptors', \
            fontsize=24, fontweight='bold')

    # ts-ss part
    ''' old
    features = {'E_CH3ab': 0.84, 'E_Hab3': 0.50, \
            'a_O2-M4-C_CH3ab': 0.41, 'd_H1-C_CH3ab': -0.32, \
            'h_O2-C-M4-H1_CH3ab': 0.30, 'h_O2-H1-M4-C_CH3ab': -0.29, \
            'a_O2-M4-H1_Hab3': -0.24, 'a_O2-C-H1_CH3ab': -0.13, \
            'a_M4-C-H1_CH3ab': 0.11, 'a_O2-H1-M4_Hab3': 0.05, \
            'd_O2-C_CH3ab': 0.02}

    feanames = ['$E_{{CH}_3^*}$', '$E_{H^*}$', \
            '$A_{O_{br}-M_{cus}-C}^{CH_3^*}$', '$D_{H-C}^{CH_3^*}$', \
            '$H_{O_{br}-C-M_{cus}-H}^{CH_3^*}$', \
            '$H_{O_{br}-H-M_{cus}-C}^{CH_3^*}$', \
            '$A_{O_{br}-M_{cus}-H}^{H^*}$', '$A_{O_{br}-C-H}^{CH_3^*}$', \
            '$A_{M_{cus}-C-H}^{CH_3^*}$', '$A_{O_{br}-H-M_{cus}}^{H^*}$', \
            '$D_{O_{br}-C}^{{CH}_3^*}$']
    '''
    features = {'E_CH3ab': 0.84, 'E_Hab3': 0.50, \
            'a_O2-M4-C_CH3ab': 0.40, 'd_H1-C_CH3ab': -0.31, \
            'h_O2-C-M4-H1_CH3ab': 0.30, 'h_O2-H1-M4-C_CH3ab': -0.29, \
            'a_O2-M4-H1_Hab3': -0.24, 'a_O2-C-H1_CH3ab': -0.12, \
            'a_M4-C-H1_CH3ab': 0.11, 'a_O2-H1-M4_Hab3': 0.05, \
            'd_O2-C_CH3ab': 0.01}

    feanames = ['$E_{{CH}_3^*}$', '$E_{H^*}$', \
            '$A_{O_{br}-M_{cus}-C}^{CH_3^*}$', '$D_{H-C}^{CH_3^*}$', \
            '$H_{O_{br}-C-M_{cus}-H}^{CH_3^*}$', \
            '$H_{O_{br}-H-M_{cus}-C}^{CH_3^*}$', \
            '$A_{O_{br}-M_{cus}-H}^{H^*}$', '$A_{O_{br}-C-H}^{CH_3^*}$', \
            '$A_{M_{cus}-C-H}^{CH_3^*}$', '$A_{O_{br}-H-M_{cus}}^{H^*}$', \
            '$D_{O_{br}-C}^{{CH}_3^*}$']

    # ts-ss ax1 features selected by lasso
    plot_selected_features(ax[0][0], features, feanames)

    # ts-ss ax2 feature selection
    plot_feature_selection(ax[0][1], 'Ets', features)
    ax[0][1].text(-2, 1.2, '(a) Surface-stabilized Mechanism', \
            fontsize=20, fontweight='normal')

    # ts-ss ax3 best linear regression
    plot_best_regression(ax[0][2], 'Ets', features, 7)

    # ts-ra part
    ''' old
    features = {'E_Hab3': 0.92, \
            'd_O2-C_CH3ab': 0.04, 'd_O2-H1_Hab3': -0.02, 'd_O2-M4_suf': 0.01, \
            'h_O2-M4-H1-C_CH3ab': 0.01}

    feanames = ['$E_{H^*}$', \
            '$D_{O_{br}-C}^{CH_3^*}$', '$D_{O_{br}-H}^{H^*}$', \
            '$D_{O_{br}-M_{cus}}^{surface}$', '$H_{O_{br}-M_{cus}-H-C}^{CH_3^*}$']
    '''
    features = {'E_Hab3': 0.86, \
            'd_O2-H1_Hab3': -0.02, 'd_O2-M4_suf': 0.01}

    feanames = ['$E_{H^*}$', \
            '$D_{O_{br}-H}^{H^*}$', \
            '$D_{O_{br}-M_{cus}}^{surface}$']

    # ts-ra ax4 features selected by lasso
    plot_selected_features(ax[1][0], features, feanames)

    # ts-ra ax5 feature selection
    plot_feature_selection(ax[1][1], 'Etsra', features)
    ax[1][1].text(0.5, 1.2, '(b) Radical-like Mechanism', \
            fontsize=20, fontweight='normal')

    # ts-ra ax6 best linear regression
    plot_best_regression(ax[1][2], 'Etsra', features, 1)

    #plt.show()
    plt.savefig('./figures/fig6.png')

if __name__ == '__main__':
    DS_CSV = '../CH4_10.csv'
    plot_main()


import csv
import json

csvfile = open('paldea.csv', 'r')
jsonfile = open('paldea-dex.json', 'w')

pokemons = []
data = {}
fieldnames = ("Number", "Status")
reader = csv.DictReader( csvfile, fieldnames)
for row in reader:
    pokemon = row
    pokemonNumber = pokemon['Number']
    newLine = {}
    newLine[pokemonNumber] = pokemon['Status']
    pokemons.append(newLine)
data["pokemon"] = pokemons
json.dump(data, jsonfile)
jsonfile.write(',\n')


from textscrub import clean, nlp
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# final function
def normalize_text(
    text,
    remove_glyphs=True,
    remove_html_tags=True,
    remove_hyperlinks=True,
    remove_punctuation=True,
    remove_spaces=True,
    tokenizing=True,
    stopwords_removal=True,
    lemmat=True,
    remove_emojis=True,
    replace_emojis=True
):
    """
    Normalize various aspects of a raw text.
    A convenience function for applying all preprocessing functions in one go.
    Args:
        text(str) -- raw text to preprocess
        remove_glyphs(bool) -- if True, remove the non-ascii characters,
        non-printable characters, non-latin characters, accents.

    Returns:
        str -- prerocessed text.
    """

    if text is None:
        return ""

    text = str(text)

    if remove_glyphs:
        text = clean.remove_glyphs(text)
    if remove_html_tags:
        text = clean.remove_html_tags(text)
    if remove_hyperlinks:
        text = clean.remove_hyperlinks(text)
    if remove_punctuation:
        text = clean.remove_punctuation(text)
    if remove_spaces:
        text = clean.remove_spaces(text)
    if stopwords_removal:
        text = nlp.stopword_removal(text)
    if lemmat:
        text = nlp.lemmat(text)
    if remove_emojis:
        text = nlp.remove_emoji(text, "remove")
    if replace_emojis:
        text = nlp.remove_emoji(text, "replace")
    if tokenizing:
        text = nlp.tokenizing(text)

    return text


import pyMeow as pm
from random import randint

def main():
    pm.overlay_init()
    fps = pm.get_monitor_refresh_rate()
    pm.set_fps(fps)
    pm.toggle_mouse()
    pm.gui_fade(0.9)
    radius = 50
    width, height = pm.get_screen_width(), pm.get_screen_height()
    x, y = width // 2, height // 2
    speed = 5
    max_speed = 100
    ball_left, ball_down = False, False
    circle_color = pm.get_color("#000000")
    stars = [pm.vec2(randint(0, width), randint(0, height)) for _ in range(300)]

    while pm.overlay_loop():
        if ball_left:
            if x > width - radius:
                ball_left = False
            else:
                x += speed
        else:
            if x < -1 + radius:
                ball_left = True
            else:
                x -= speed

        if ball_down:
            if y > height - radius:
                ball_down = False
            else:
                y += speed
        else:
            if y < -1 + radius:
                ball_down = True
            else:
                y -= speed

        pm.begin_drawing()
        [pm.draw_pixel(vec["x"], vec["y"], pm.get_color("#ffffff")) for vec in stars]
        pm.draw_circle(x, y, radius, circle_color)
        pm.draw_fps(width / 2, height / 2)

        if pm.gui_button(0, 0, 100, 50, "Increase speed"):
            if speed != max_speed:
                speed += 5
        elif pm.gui_button(0, 50, 100, 50, "Decrease speed"):
            if speed != 0:
                speed -= 5
        elif pm.gui_button(100, 0, 100, 50, "Change Color"):
            circle_color = pm.new_color(
                randint(0, 255),
                randint(0, 255),
                randint(0, 255),
                randint(0, 255)
            )
        elif pm.gui_button(100, 50, 100, 50, "Mouse passthrough"):
            pm.toggle_mouse()
        elif pm.gui_button(200, 0, 100, 50, "Increase FPS"):
            fps += 1
            pm.set_fps(fps)
        elif pm.gui_button(200, 50, 100, 50, "Decrease FPS"):
            fps -= 1
            pm.set_fps(fps)

        pm.gui_progress_bar(350, 10, 200, 30, "Speed: ", str(speed), speed, 0, max_speed)
        radius = pm.gui_slider(350, 60, 200, 30, "Radius", str(radius), radius, 0, 150)
        pm.draw_text("Exit with 'END'", 50, 100, 25, pm.get_color("#ffffff"))
        pm.end_drawing()

if __name__ == "__main__":
    main()


if __name__ == "__main__":
    import sqlite3

    connection_obj = sqlite3.connect('db-cliente.db')

    cursor_obj = connection_obj.cursor()

    cursor_obj.execute("CREATE TABLE clientes (  \
        id              INTEGER PRIMARY KEY AUTOINCREMENT, \
        email           TEXT, \
        nome	        TEXT,   \
        cpf	            NUMERIC, \
        telefone        NUMERIC, \
        sexo	        CHARACTER(1), \
        data_nascimento	DATE, \
        cep	            NUMERIC, \
        endereco	    TEXT, \
        numero	        NUMERIC, \
        bairro	        TEXT \
    );")
    connection_obj.commit()


import ROOT
from ROOT import *
import math
from array import array

import argparse
parser =  argparse.ArgumentParser(description='Photon ID MVA Optimizer')
parser.add_argument('-s', '--signal', dest='sig', required=True, type=str)
parser.add_argument('-b', '--background', dest='bac', required=True, type=str)
parser.add_argument('-v', '--var', dest='var', required=True, type=str)
parser.add_argument('-c', '--cut', dest='cut', required=True, type=str)
parser.add_argument('-n', '--name', dest='name', required=True, type=str)
parser.add_argument('-o', '--outputLoc', dest='outputLoc', required=True, type=str)

opt = parser.parse_args()
sigMass = opt.sig
bkg = opt.bac
var = opt.var
cut = opt.cut
outputLoc = opt.outputLoc
name = opt.name

# print opt.cut
f_sig = TFile('/eos/user/t/twamorka/newCatalog_fixVtx_3Oct2019/hadd_Tree/signal_m_'+str(sigMass)+'.root')
t_sig = f_sig.Get('h4gCandidateDumper/trees/SUSYGluGluToHToAA_AToGG_M_'+str(sigMass)+'_TuneCUETP8M1_13TeV_pythia8_13TeV_4photons')

bkgFiles = []
if bkg == 'QCD':
   bkgFiles.append(['/eos/user/t/twamorka/2016Ntuples/QCD30to40.root','QCD_Pt_30to40_DoubleEMEnriched_MGG_80toInf_TuneCUETP8M1_13TeV_Pythia8_13TeV_4photons','QCD 30to40'])
   bkgFiles.append(['/eos/user/t/twamorka/2016Ntuples/QCD30toInf.root','QCD_Pt_30toInf_DoubleEMEnriched_MGG_40to80_TuneCUETP8M1_13TeV_Pythia8_13TeV_4photons','QCD 30toInf'])
   bkgFiles.append(['/eos/user/t/twamorka/2016Ntuples/QCD40toInf.root','QCD_Pt_40toInf_DoubleEMEnriched_MGG_80toInf_TuneCUETP8M1_13TeV_Pythia8_13TeV_4photons','QCD 40toInf'])

elif bkg == 'GJet':
   bkgFiles.append(['/eos/user/t/twamorka/2016Ntuples/GJet20to40.root','GJet_Pt_20to40_DoubleEMEnriched_MGG_80toInf_TuneCUETP8M1_13TeV_Pythia8_13TeV_4photons','GJet 20to40'])
   bkgFiles.append(['/eos/user/t/twamorka/2016Ntuples/GJet20toInf.root','GJet_Pt_20toInf_DoubleEMEnriched_MGG_40to80_TuneCUETP8M1_13TeV_Pythia8_13TeV_4photons','GJet 20toInf'])
   bkgFiles.append(['/eos/user/t/twamorka/2016Ntuples/GJet40toInf.root','GJet_Pt_40toInf_DoubleEMEnriched_MGG_80toInf_TuneCUETP8M1_13TeV_Pythia8_13TeV_4photons','GJet 40toInf'])

elif bkg == 'DiPho':
   bkgFiles.append(['/eos/user/t/twamorka/2016Ntuples/DiPho40to80.root','DiPhotonJetsBox_M40_80_Sherpa_13TeV_4photons','DiPho 40to80'])
   bkgFiles.append(['/eos/user/t/twamorka/2016Ntuples/DiPho80toInf.root','DiPhotonJetsBox_MGG_80toInf_13TeV_Sherpa_13TeV_4photons','DiPho 80toInf'])

elif bkg == 'All':
   bkgFiles.append(['/eos/user/t/twamorka/2016Ntuples/QCD30to40.root','QCD_Pt_30to40_DoubleEMEnriched_MGG_80toInf_TuneCUETP8M1_13TeV_Pythia8_13TeV_4photons','QCD 30to40'])
   bkgFiles.append(['/eos/user/t/twamorka/2016Ntuples/QCD30toInf.root','QCD_Pt_30toInf_DoubleEMEnriched_MGG_40to80_TuneCUETP8M1_13TeV_Pythia8_13TeV_4photons','QCD 30toInf'])
   bkgFiles.append(['/eos/user/t/twamorka/2016Ntuples/QCD40toInf.root','QCD_Pt_40toInf_DoubleEMEnriched_MGG_80toInf_TuneCUETP8M1_13TeV_Pythia8_13TeV_4photons','QCD 40toInf'])
   bkgFiles.append(['/eos/user/t/twamorka/2016Ntuples/GJet20to40.root','GJet_Pt_20to40_DoubleEMEnriched_MGG_80toInf_TuneCUETP8M1_13TeV_Pythia8_13TeV_4photons','GJet 20to40'])
   bkgFiles.append(['/eos/user/t/twamorka/2016Ntuples/GJet20toInf.root','GJet_Pt_20toInf_DoubleEMEnriched_MGG_40to80_TuneCUETP8M1_13TeV_Pythia8_13TeV_4photons','GJet 20toInf'])
   bkgFiles.append(['/eos/user/t/twamorka/2016Ntuples/GJet40toInf.root','GJet_Pt_40toInf_DoubleEMEnriched_MGG_80toInf_TuneCUETP8M1_13TeV_Pythia8_13TeV_4photons','GJet 40toInf'])
   bkgFiles.append(['/eos/user/t/twamorka/2016Ntuples/DiPho40to80.root','DiPhotonJetsBox_M40_80_Sherpa_13TeV_4photons','DiPho 40to80'])
   bkgFiles.append(['/eos/user/t/twamorka/2016Ntuples/DiPho80toInf.root','DiPhotonJetsBox_MGG_80toInf_13TeV_Sherpa_13TeV_4photons','DiPho 80toInf'])

Cut_Cat = []
i = -0.9
while i<1:
    Cut_Cat.append(['pho1_MVA > '+str(i)+' && pho2_MVA >'+str(i)+'&&'+cut, 'pho1_MVA > '+str(i)+' && pho2_MVA >'+str(i)+' && pho3_MVA >' + str(i) +' && pho4_MVA >'+str(i)+'&&'+cut,'pho1_MVA > -0.9 && pho2_MVA > -0.9'+' && pho3_MVA >'+str(i)+' && pho4_MVA >'+str(i)+'&&'+cut] )
    i += 0.1

# print Cut_Cat
t_bkg = TChain()
# print bkgFiles
for file in bkgFiles:
    t_bkg.Add(str(file[0])+'/h4gCandidateDumper/trees/'+str(file[1]))

# print t_bkg.GetEntries()

t_sig.Draw(var + '>> h_sig')
t_bkg.Draw(var + '>> h_bkg')
tot_sig =  h_sig.Integral()
tot_bkg = h_bkg.Integral()
hists = []
graphs  = TMultiGraph()
sigEff_Cat0, bkgEff_Cat0 = array( 'd' ), array( 'd' )
sigEff_Cat1, bkgEff_Cat1 = array( 'd' ), array( 'd' )
sigEff_Cat2, bkgEff_Cat2 = array( 'd' ), array( 'd' )
for cut in Cut_Cat:
    t_sig.Draw(var + '>> h0_sig',TCut(cut[0]))
    t_bkg.Draw(var + '>> h0_bkg',TCut(cut[0]))
    sigEff_Cat0.append(h0_sig.Integral()/tot_sig)
    bkgEff_Cat0.append(1-(h0_bkg.Integral()/tot_bkg))

    t_sig.Draw(var + '>> h1_sig',TCut(cut[1]))
    t_bkg.Draw(var + '>> h1_bkg',TCut(cut[1]))
    sigEff_Cat1.append(h1_sig.Integral()/tot_sig)
    bkgEff_Cat1.append(1-(h1_bkg.Integral()/tot_bkg))

    t_sig.Draw(var + '>> h2_sig',TCut(cut[2]))
    t_bkg.Draw(var + '>> h2_bkg',TCut(cut[2]))
    sigEff_Cat2.append(h2_sig.Integral()/tot_sig)
    bkgEff_Cat2.append(1-(h2_bkg.Integral()/tot_bkg))

# print bkgEff_Cat0
gr1 = TGraph(len(sigEff_Cat0),sigEff_Cat0,bkgEff_Cat0)
gr1.SetMarkerColor(kBlack)
gr1.SetMarkerStyle(20)
gr1.SetMarkerSize(1.5)

gr2 = TGraph(len(sigEff_Cat1),sigEff_Cat1,bkgEff_Cat1)
gr2.SetMarkerColor(kRed)
gr2.SetMarkerStyle(20)
gr2.SetMarkerSize(1.5)

gr3 = TGraph(len(sigEff_Cat2),sigEff_Cat2,bkgEff_Cat2)
gr3.SetMarkerColor(kGreen)
gr3.SetMarkerStyle(20)
gr3.SetMarkerSize(1.5)

c0 = TCanvas("c", "c", 800, 750)
c0.SetGridy()
c0.SetGridx()
leg = TLegend(0.152882, 0.167143, 0.442356, 0.357143)
leg.SetBorderSize(0)
leg.SetTextSize(0.04)
leg.SetFillColor(kWhite)
leg.SetFillStyle(0)
graphs.Add(gr1)
graphs.Add(gr2)
graphs.Add(gr3)
leg.AddEntry(gr1,"Cut 1","lp")
leg.AddEntry(gr2,"Cut 2","lp")
leg.AddEntry(gr3,"Cut 3","lp")
graphs.SetTitle(";Signal Efficiency; 1- Background Efficiency")
graphs.Draw("AP")
leg.Draw("same")
c0.SaveAs(outputLoc+'Eff_'+str(bkg)+'_'+str(opt.name)+'.pdf')
c0.SaveAs(outputLoc+'Eff_'+str(bkg)+'_'+str(opt.name)+'.png')


import pandas as pd
import matplotlib.pyplot as plt

df_tesla = pd.read_csv('TESLA Search Trend vs Price.csv')

df_btc_search = pd.read_csv('Bitcoin Search Trend.csv')
df_btc_price = pd.read_csv('Daily Bitcoin Price.csv')

df_unemployment = pd.read_csv('UE Benefits Search vs UE Rate 2004-19.csv')

# Challenge:

# 1.What are the shapes of the dataframes?
# 2. How many rows and columns?
# 3. What are the column names?
# 4. Complete the f-string to show the largest/smallest number in the search data column
# 5. Try the .describe() function to see some useful descriptive statistics
# 6. What is the periodicity of the time series data (daily, weekly, monthly)?
# 7. What does a value of 100 in the Google Trend search popularity actually mean?

# Tesla
# 1, 2, 3
print(df_tesla.shape)
df_tesla.head()

# 4
print(f'Largest value for Tesla in Web Search: {df_tesla["TSLA_WEB_SEARCH"].max()}')
print(f'Smallest value for Tesla in Web Search: {df_tesla["TSLA_WEB_SEARCH"].min()}')

# 5
df_tesla.describe()

# Unemployment

df_unemployment.shape
df_unemployment.head()

print('Largest value for "Unemployemnt Benefits" '
      f'in Web Search: {df_unemployment["UE_BENEFITS_WEB_SEARCH"].max()}')

#Bitcoin

df_btc_price.shape
df_btc_price.head()

df_btc_search.shape
df_btc_search.head()

print(f'largest BTC News Search: {df_btc_search["BTC_NEWS_SEARCH"].max()}')

# Challenge: Are there any missing values in any of the dataframes? If so,
# which row/rows have missing values? How many missing values are there?

print(f'Missing values for Tesla?: {df_tesla.isna().values.any()}')
print(f'Missing values for U/E?: {df_unemployment.isna().values.any()}')
print(f'Missing values for BTC Search?: {df_btc_search.isna().values.any()}')

# This one has some
print(f'Missing values for BTC price?: {df_btc_price.isna().values.any()}')
# Two to be exact
print(f'Number of missing values: {df_btc_price.isna().values.sum()}')

# Challenge: Remove any missing values that you found.
df_btc_price = df_btc_price.dropna()

# Challenge
# Our DataFrames contain time-series data.
# Do you remember how to check the data type of the entries in the DataFrame?
# Have a look at the data types of the MONTH or DATE columns.
# Convert any strings you find into Datetime objects.
# Do this for all 4 DataFrames. Double-check if your type conversion was successful.

df_tesla.dtypes

df_tesla.MONTH = pd.to_datetime(df_tesla.MONTH)
df_btc_search.MONTH = pd.to_datetime(df_btc_search.MONTH)
df_unemployment.MONTH = pd.to_datetime(df_unemployment.MONTH)
df_btc_price.DATE = pd.to_datetime(df_btc_price.DATE)

# Challenge
# Plot the Tesla stock price against the Tesla search volume using a line chart and
# two different axes.

ax1 = plt.gca()
ax2 = ax1.twinx()

ax1.set_ylabel('TSLA Stock Price')
ax2.set_ylabel('Search Trend')

ax1.plot(df_tesla.MONTH, df_tesla.TSLA_USD_CLOSE)
ax2.plot(df_tesla.MONTH, df_tesla.TSLA_WEB_SEARCH)

# Challenge
# Try using one of the blue colour names for the search volume and
#  a HEX code for a red colour for the stock price.

ax1 = plt.gca()
ax2 = ax1.twinx()

ax1.set_ylabel('TSLA Stock Price', color='#E6232E') # can use a HEX code
ax2.set_ylabel('Search Trend', color='skyblue') # or a named colour

ax1.plot(df_tesla.MONTH, df_tesla.TSLA_USD_CLOSE, color='#E6232E')
ax2.plot(df_tesla.MONTH, df_tesla.TSLA_WEB_SEARCH, color='skyblue')

# Challenge
# There are still some ways to improve the look of this chart. First off, let's make it larger. Can you make the following changes:

# Increase the figure size (e.g., to 14 by 8).

# Increase the font sizes for the labels and the ticks on the x-axis to 14.

# Rotate the text on the x-axis by 45 degrees.

# Add a title that reads 'Tesla Web Search vs Price'

# Make the lines on the chart thicker.

# Keep the chart looking sharp by changing the dots-per-inch or DPI value.

# Set minimum and maximum values for the y and x-axis. Hint: check out methods like set_xlim().

# Finally use plt.show() to display the chart below the cell instead of relying on the automatic notebook output.

plt.figure(figsize=(14,8), dpi=120)
plt.title('Tesla Web Search vs Price', fontsize=18)

# Increase the size and rotate the labels on the x-axis
plt.xticks(fontsize=14, rotation=45)

ax1 = plt.gca()
ax2 = ax1.twinx()

ax1.set_ylabel('TSLA Stock Price', color='#E6232E', fontsize=14)
ax2.set_ylabel('Search Trend', color='skyblue', fontsize=14)

# Set the minimum and maximum values on the axes
ax1.set_ylim([0, 600])
ax1.set_xlim([df_tesla.MONTH.min(), df_tesla.MONTH.max()])

ax1.plot(df_tesla.MONTH, df_tesla.TSLA_USD_CLOSE, color='#E6232E', linewidth=3)
ax2.plot(df_tesla.MONTH, df_tesla.TSLA_WEB_SEARCH, color='skyblue', linewidth=3)

plt.show()

# Challenge
# Plot the search for "unemployment benefits" against the official unemployment rate.

# Change the title to: Monthly Search of "Unemployment Benefits" in the U.S. vs the U/E Rate

# Change the y-axis label to: FRED U/E Rate

# Change the axis limits

# Add a grey grid to the chart to better see the years and the U/E rate values. Use dashed lines for the line style.

# Can you discern any seasonality in the searches? Is there a pattern?

plt.figure(figsize=(14,8), dpi=120)
plt.title('Monthly Search of "Unemployment Benefits" in the U.S. vs the U/E Rate', fontsize=18)
plt.yticks(fontsize=14)
plt.xticks(fontsize=14, rotation=45)

ax1 = plt.gca()
ax2 = ax1.twinx()

ax1.set_ylabel('FRED U/E Rate', color='purple', fontsize=14)
ax2.set_ylabel('Search Trend', color='skyblue', fontsize=14)

# ax1.xaxis.set_major_locator(years)
# ax1.xaxis.set_major_formatter(years_fmt)
# ax1.xaxis.set_minor_locator(months)

ax1.set_ylim(bottom=3, top=10.5)
ax1.set_xlim([df_unemployment.MONTH.min(), df_unemployment.MONTH.max()])

# Show the grid lines as dark grey lines
ax1.grid(color='grey', linestyle='--')

# Change the dataset used
ax1.plot(df_unemployment.MONTH, df_unemployment.UNRATE,
         color='purple', linewidth=3, linestyle='--')
ax2.plot(df_unemployment.MONTH, df_unemployment.UE_BENEFITS_WEB_SEARCH,
         color='skyblue', linewidth=3)

plt.show()

Advanced/Day75/Google Trends Data Viz (start)/main.py


"""
This is a better implementation that uses a decorator and prevents the initializer
from being called multiple times.
"""

def singleton(class_):
    _instances = {}

    def get_instance(*args, **kwargs):
        if class_ not in _instances:
            _instances[class_] = class_(*args, **kwargs)
        return _instances[class_]

    return get_instance

@singleton
class Database:
    def __init__(self):
        print("Loading database...")

import unittest

class SingletonPattern(unittest.TestCase):
    def test_singleton_pattern(self):
        d1 = Database()
        d2 = Database()
        print("Initializer called once")

        self.assertTrue(d1 == d2)

if __name__ == "__main__":
    unittest.main()


import json

from django.http       import JsonResponse
from django.views      import View

from postings.models   import Post, Image
from users.models      import User
from core.utils        import login_decorator

class PostingView(View) :
    @login_decorator
    def post(self, request) :
        try :
            data     = json.loads(request.body)
            user     = request.user
            content    = data['content']
            image_list = data['image'].split(',')
            post = Post.objects.create(
                content = content,
                user = user
            )

            for image in image_list :
                Image.objects.create(
                    url = image,
                    post = post
                )

            return JsonResponse({'message':'SUCCESS'}, status=200)

        except KeyError :
            return JsonResponse({'message':'KEY_ERROR'}, status=400)

    @login_decorator
    def get(self, request) :
        post_list = [{
            'user'       : User.objects.get(id=post.user.id).id,
            'content'    : post.content,
            'images'     : [image.url for image in post.images.all()],
            'created_at' : post.created_at
        } for post in Post.objects.all()
        ]

        return JsonResponse({'data' : post_list}, status=200)


import smbus
import time
bus = smbus.SMBus(1)
address = 0x70

#SRF08 REQUIRES 5V
time.sleep(1)

def lightlevel():
        light = bus.read_byte_data(address, 1)
        return light

def brightness(light):
    if light <50:
        print("Too Dark")
    if light >50 and light <100:
        print("Dark")
    if light >100 and light <150:
        print("Medium")
    if light >150 and light <200:
        print("Bright")
    if light >200 and light <250:
        print("Too Bright")

while True:

        time.sleep(0.7)
        lightlvl = lightlevel()

        print(brightness(lightlvl))


from env_FindGoals import EnvFindGoals
from GDICE_MP import GDICEopt, is_obs_in_list
import random
import matplotlib.pyplot as plt

if __name__ == '__main__':
    # running GDICE
    alpha = 0.05
    N_agent = 2
    N_s = 20
    N_b = 5     # no bigger than N_s
    N_action = 5
    N_node_list = [10, 10]
    max_opt_iter = 100
    max_MC_iter = 300

    # get obs_dict
    obs_list1 = []
    obs_list2 = []
    env = EnvFindGoals()
    print('start')
    for j in range(100):
        print(j)
        env.reset()
        for i in range(100):
            reward_list, done = env.step([random.randint(0, 4), random.randint(0, 4)])
            obs1 = env.get_obs()[0]
            obs2 = env.get_obs()[1]
            if is_obs_in_list(obs1, obs_list1) == False:
                obs_list1.append(obs1)
            if is_obs_in_list(obs2, obs_list2) == False:
                obs_list2.append(obs2)
        print(len(obs_list1))
        print(len(obs_list2))
    obs_dict = [obs_list1, obs_list2]
    print('here')

    #obs_dict = [[0], [0], [0]]
    env = EnvFindGoals()
    optimizer = GDICEopt(N_agent, N_node_list, obs_dict, N_action, N_s, N_b, max_opt_iter, max_MC_iter, env)
    #optimizer.load_model('GDICE_model_FindGoals')

    optimizer.optimize()
    best_ctrl = optimizer.best_ctrl
    # optimizer.save_train_curve()
    env.reset()
    best_ctrl.reset_state()
    for i in range(max_MC_iter):
        action_list = best_ctrl.get_action()
        reward_list, done = env.step(action_list)
        obs_list = env.get_obs()

        best_ctrl.state_transit(obs_list)
        if done:
            best_ctrl.reset_state()
            env.reset()
        env.plot_scene()

        '''# save global plot
        plt.figure()
        full_obs = env.get_full_obs()
        plt.imshow(full_obs)
        plt.savefig('/home/shuo/Figures/' + str(i) + '.png')
        plt.close()'''


"""add last notice and internal addresses

Revision ID: 2e6851ff437
Revises: 42a3f495ec4
Create Date: 2014-06-05 21:40:47.132595

"""

# revision identifiers, used by Alembic.
revision = '2e6851ff437'
down_revision = '42a3f495ec4'

from alembic import op
import sqlalchemy as sa

def upgrade():
    ### commands auto generated by Alembic - please adjust! ###
    op.add_column('users', sa.Column('last_expiry_notice', sa.DateTime(), nullable=True))
    op.add_column('vpnsessions', sa.Column('internal_ip4', sa.String(), nullable=True))
    op.add_column('vpnsessions', sa.Column('internal_ip6', sa.String(), nullable=True))
    ### end Alembic commands ###

def downgrade():
    ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('vpnsessions', 'internal_ip6')
    op.drop_column('vpnsessions', 'internal_ip4')
    op.drop_column('users', 'last_expiry_notice')
    ### end Alembic commands ###


#!/usr/bin/env python3

import requests
import shelve
from operator import itemgetter
from lxml import etree
import sys

us=True

def getData():
  data=[]
  s=1
  maxS=1

  while s<=maxS:
    page=1
    maxPage=1
    while page<=maxPage:
      print('Now processing %dS, page: %d' % (s, page))
      r=requests.get('http://www.hobbyking.com/hobbyking/store/uh_listCategoriesAndProducts.asp?cwhl=%s&pc=&idCategory=86&curPage=%d&v=&sortlist=%s&sortMotor=&LiPoConfig=%d&CatSortOrder=desc' % (('US' if us else 'XX'), page, ('H' if us else 'XX'), s))
      tree=etree.HTML(r.content)
      for input in tree.xpath("//div[@id='tableContent']/div/table/tr/td/table/tr/td/input"):
        sInput=int(input.get('onclick').split('LiPoConfig=')[1].split('&')[0])
        if sInput>maxS:
          maxS=sInput
      for link in tree.xpath("//div[@id='tableContent']/div/form/ul/li/a"):
        pageLink=int(link.get('href').split('curPage=')[1].split('&')[0])
        if pageLink>maxPage:
          maxPage=pageLink
      for item in tree.xpath("//div[@id='tableContent']/div/table/tr[@class='zeroLineHeight']/td/table/tr/td/table/tr/td/div/div/div/span/strong/font[contains(.,'$')]"):
        price=float(item.text.strip().lstrip('$'))
        itemA=item.xpath('../../../../../../div/a')
        href=itemA[0].get('href')
        mAh=[float(x.lower().rstrip('mah')) for x in itemA[0].text.split(' ') if x.lower().endswith('mah')]
        if len(mAh)>0:
          mAh=mAh[0]
        else:
          mAh=[float(x) for x in itemA[0].text.split(' ') if x.isdigit()]
          if len(mAh)>0:
            if len(mAh)>1:
              print('Found more than one mAh for %s, using the first one' % itemA[0].text)
            mAh=mAh[0]
          else:
            print("Couldn't find a mAh for %s" % itemA[0].text)
            continue
        Wh=s*3.7*mAh/1000
        data.append([Wh/price, price, Wh, mAh, s, href, itemA[0].text])
      page+=1
    s+=1
  d=shelve.open('batteries-W%s-V%s.data' % (('US' if us else 'XX'), str(sys.version_info[0])))
  d['data']=data
  d.close()

def processData():
  d=shelve.open('batteries-W%s-V%s.data' % (('US' if us else 'XX'), str(sys.version_info[0])))
  data=d['data']
  d.close()
  data.sort(key=itemgetter(0), reverse=True)
  for item in data[:20]:
    print('Wh/$: %4.2f, Wh: %5.2f, $: %5.2f, s: %s, %s, %s' % (item[0], item[2], item[1], item[4], item[6], item[5]))

getData()
processData()


import ctypes  # type: ignore
import numpy as np  # type: ignore
import os

def load_library(name):
	for prefix in [""] + os.environ["PATH"].split(os.pathsep):
		try:
			# print(os.path.join(prefix, name))
			return ctypes.cdll.LoadLibrary(os.path.join(prefix, name))
		except OSError:
			# print("    ", err)
			pass

		try:
			# print(os.path.join(prefix, "lib" + name + ".so"))
			return ctypes.cdll.LoadLibrary(
			    os.path.join(prefix, "lib" + name + ".so"))
		except OSError:
			# print("    ", err)
			pass
	raise OSError(name + " not found.")

minimum_bibliotek = load_library("minimum_bibliotek")

minimum_bibliotek_error_string_ptr = ctypes.c_char_p.in_dll(
    minimum_bibliotek, "error_string")

def error_string() -> str:
	val = minimum_bibliotek_error_string_ptr.value
	if val:
		return val.decode("utf-8")
	else:
		return ""

ctypes_minimum_linear_test = minimum_bibliotek.minimum_linear_test
ctypes_minimum_linear_test.restype = ctypes.c_int

def minimum_linear_test(data):
	assert data.dtype == np.int32
	return ctypes_minimum_linear_test(
	    data.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
	    ctypes.c_int(data.size))

class Ip_p(ctypes.c_void_p):
	# http://stackoverflow.com/questions/17840144/why-does-setting-ctypes-dll-function-restype-c-void-p-return-long
	pass

class Solver_p(ctypes.c_void_p):
	pass

class Solutions_p(ctypes.c_void_p):
	pass

ctypes_minimum_linear_ip = minimum_bibliotek.minimum_linear_ip
ctypes_minimum_linear_ip.restype = Ip_p

def minimum_linear_ip(proto_byte_string: bytes) -> Ip_p:
	return ctypes_minimum_linear_ip(
	    ctypes.c_char_p(proto_byte_string),
	    ctypes.c_int32(len(proto_byte_string)))

ctypes_minimum_linear_ip_destroy = minimum_bibliotek.minimum_linear_ip_destroy
ctypes_minimum_linear_ip_destroy.restype = None

def minimum_linear_ip_destroy(ip: Ip_p):
	assert type(ip) == Ip_p
	return ctypes_minimum_linear_ip_destroy(ip)

ctypes_minimum_linear_ip_save_mps = minimum_bibliotek.minimum_linear_ip_save_mps
ctypes_minimum_linear_ip_save_mps.restype = None

def minimum_linear_ip_save_mps(ip: Ip_p, file_name: str):
	assert type(ip) == Ip_p
	assert type(file_name) == str
	ctypes_minimum_linear_ip_save_mps(
	    ip, ctypes.c_char_p(file_name.encode("utf-8")))

ctypes_minimum_linear_solver_default = minimum_bibliotek.minimum_linear_solver_default
ctypes_minimum_linear_solver_default.restype = Solver_p

def minimum_linear_solver_default():
	return ctypes_minimum_linear_solver_default()

ctypes_minimum_linear_solver_minisat = minimum_bibliotek.minimum_linear_solver_minisat
ctypes_minimum_linear_solver_minisat.restype = Solver_p

def minimum_linear_solver_minisat():
	return ctypes_minimum_linear_solver_minisat()

ctypes_minimum_linear_solver_destroy = minimum_bibliotek.minimum_linear_solver_destroy
ctypes_minimum_linear_solver_destroy.restype = None

def minimum_linear_solver_destroy(solver: Solver_p):
	assert type(solver) == Solver_p
	return ctypes_minimum_linear_solver_destroy(solver)

ctypes_minimum_linear_solutions = minimum_bibliotek.minimum_linear_solutions
ctypes_minimum_linear_solutions.restype = Solutions_p

def minimum_linear_solutions(solver: Solver_p, ip: Ip_p):
	assert type(ip) == Ip_p
	assert type(solver) == Solver_p
	solutions = ctypes_minimum_linear_solutions(solver, ip)
	assert solutions, error_string()
	return solutions

ctypes_minimum_linear_solutions_get = minimum_bibliotek.minimum_linear_solutions_get
ctypes_minimum_linear_solutions_get.restype = ctypes.c_int

def minimum_linear_solutions_get(solutions: Solutions_p, ip: Ip_p, solution):
	assert type(solutions) == Solutions_p
	assert type(ip) == Ip_p
	assert solution.dtype == np.float64
	return ctypes_minimum_linear_solutions_get(
	    solutions, ip, solution.ctypes.data_as(
	        ctypes.POINTER(ctypes.c_double)), ctypes.c_int32(solution.size))

ctypes_minimum_linear_solutions_destroy = minimum_bibliotek.minimum_linear_solutions_destroy
ctypes_minimum_linear_solutions_destroy.restype = None

def minimum_linear_solutions_destroy(solutions: Solutions_p):
	assert type(solutions) == Solutions_p
	return ctypes_minimum_linear_solutions_destroy(solutions)

ctypes_sum_function = minimum_bibliotek.sum_function

def sum_function(data):
	return ctypes_sum_function(
	    data.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
	    ctypes.c_int32(data.size))

ctypes_add_function = minimum_bibliotek.add_function

def add_function(a, b):
	return ctypes_add_function(ctypes.c_int32(a), ctypes.c_int32(b))

# Run a quick test to make sure everything is OK.
_a = np.array([1, 2, 3, 4, 5, 6], dtype=np.int32)
assert minimum_linear_test(_a) == 21


def part1():
    octopi = parse_input()
    flashes = 0
    steps = 100
    for step in range(0, steps):
        for y, row in enumerate(octopi):
            for x, level in enumerate(row):
                increment(octopi, y, x)

        for y, row in enumerate(octopi):
            for x, level in enumerate(row):
                if level > 9:
                    flashes += flash(octopi, y, x)
    return flashes

def part2():
    octopi = parse_input()
    step = 1
    while True:
        flashes_in_step = 0
        for y, row in enumerate(octopi):
            for x, level in enumerate(row):
                increment(octopi, y, x)

        for y, row in enumerate(octopi):
            for x, level in enumerate(row):
                if level > 9:
                    flashes_in_step += flash(octopi, y, x)

        if flashes_in_step == 100:
            return step

        step += 1

def parse_input():
    octopi = []
    with open("input.txt", "r") as input:
        for line in input:
            octopi.append([int(level) for level in list(line.strip())])
    return octopi

def increment(octopi, y, x):
    octopi[y][x] = octopi[y][x] + 1

def flash(octopi, y, x):
    flashes = 1
    octopi[y][x] = 0
    top = y - 1
    right = x + 1
    bottom = y + 1
    left = x - 1
    top_available = top >= 0
    right_available = right <= 9
    bottom_available = bottom <= 9
    left_available = left >= 0

    # top
    if top_available and octopi[top][x] != 0:
        increment(octopi, top, x)
        if (octopi[top][x] > 9):
            flashes += flash(octopi, top, x)

    # top right
    if top_available and right_available and octopi[top][right] != 0:
        increment(octopi, top, right)
        if (octopi[top][right] > 9):
            flashes += flash(octopi, top, right)

    # right
    if right_available and octopi[y][right] != 0:
        increment(octopi, y, right)
        if (octopi[y][right] > 9):
            flashes += flash(octopi, y, right)

    # bottom right
    if bottom_available and right_available and octopi[bottom][right] != 0:
        increment(octopi, bottom, right)
        if (octopi[bottom][right] > 9):
            flashes += flash(octopi, bottom, right)

    # bottom
    if bottom_available and octopi[bottom][x] != 0:
        increment(octopi, bottom, x)
        if (octopi[bottom][x] > 9):
            flashes += flash(octopi, bottom, x)

    # bottom left
    if bottom_available and left_available and octopi[bottom][left] != 0:
        increment(octopi, bottom, left)
        if (octopi[bottom][left] > 9):
            flashes += flash(octopi, bottom, left)

    # left
    if left_available and octopi[y][left] != 0:
        increment(octopi, y, left)
        if (octopi[y][left] > 9):
            flashes += flash(octopi, y, left)

    # top left
    if top_available and left_available and octopi[top][left] != 0:
        increment(octopi, top, left)
        if (octopi[top][left] > 9):
            flashes += flash(octopi, top, left)

    return flashes

if __name__ == "__main__":
    print("Part 1: " + str(part1()))
    print("Part 2: " + str(part2()))


MAIN_DATA_HUB_ID = 0

LEG_CONTROLLER_ID = 10
HEAD_SENSORS_ID = 11
TAIL_BOARD_ID = 12
LEG_SENSORS_ID = 13

class AntDeviceConfig():

    MAIN_DATA_HUB = {"DeviceId" : MAIN_DATA_HUB_ID,
                     "DeviceName" : "Main Data Hub"}

    HEAD_SENSORS = {"DeviceId" : HEAD_SENSORS_ID,
                     "DeviceName" : "head sensors"}

    LEG_CONTROLLER = {"DeviceId" : LEG_CONTROLLER_ID,
                     "DeviceName" : "Motion Controller"}

    LEG_SENSORS = {"DeviceId" : LEG_SENSORS_ID,
                     "DeviceName" : "leg sensors"}

FRONT_LEFT_LEG_LED = {  # FRONT_LEFT_LEG_LED ("front left",1),
    "name": "front left",
    "local_id": 1,
    "global_id": 0,
}

FRONT_RIGHT_LEG_LED = {  # FRONT_RIGHT_LEG_LED ("front right",4),
    "name": "front right",
    "local_id": 4,
    "global_id": 0,
}

CENTER_LEFT_LEG_LED = {  # CENTER_LEFT_LEG_LED ("center left",0),
    "name": "center left",
    "local_id": 0,
    "global_id": 0,
}

CENTER_RIGHT_LEG_LED = {  # CENTER_RIGHT_LEG_LED ("center right",3),
    "name": "center right",
    "local_id": 3,
    "global_id": 0,
}

BACK_LEFT_LEG_LED = {  # BACK_LEFT_LEG_LED ("back left",2),
    "name": "back left",
    "local_id": 2,
    "global_id": 0,
}

BACK_RIGHT_LEG_LED = {  # BACK_RIGHT_LED ("cack right",5),
    "name": "back right",
    "local_id": 5,
    "global_id": 0,
}

LEFT_LIGHT_SENSOR = {  # BACK_RIGHT_LED ("cack right",5),
    "name": "back right",
    "local_id": 0,
    "global_id": 0,
}

CENTER_LIGHT_SENSOR = {  # BACK_RIGHT_LED ("cack right",5),
    "name": "back right",
    "local_id": 1,
    "global_id": 0,
}

RIGHT_LIGHT_SENSOR = {  # BACK_RIGHT_LED ("cack right",5),
    "name": "back right",
    "local_id": 2,
    "global_id": 0,
}

class AntComponents:
    def __init__(self):
        pass

class AntComponent:
    def __init__(self, data: dict):
        self._data = data

    @property
    def local_id(self) -> int:
        return self._data["local_id"]

    @property
    def name(self) -> int:
        return self._data["name"]

    @property
    def device_id(self) -> int:
        return self._data["device_id"]

    @property
    def global_id(self) -> int:
        return self._data["global_id"]


import unittest
import update_dns

def get_event(ip4, ip6):
    params = {}
    if ip4 is not None:
        params['ip4'] = ip4
    if ip6 is not None:
        params['ip6'] = ip6
    return {'queryStringParameters': params}

class GetIpsFromEventTests(unittest.TestCase):

    def __init__(self, method_name: str = ...) -> None:
        super().__init__(method_name)

    def test_empty_ips(self):
        (ip4, ip6) = update_dns.get_ips_from_event(get_event(None, None))
        self.assertTrue(ip4 is None)
        self.assertTrue(ip6 is None)

    def test_valid_ips(self):
        (ip4, ip6) = update_dns.get_ips_from_event(
            get_event('127.0.0.1', 'ff01::4'))
        self.assertEqual(ip4, '127.0.0.1')
        self.assertEqual(ip6, 'ff01::4')

    def test_invalid_ips(self):
        (ip4, ip6) = update_dns.get_ips_from_event(
            get_event('277.0.0.1', 'ff..01::4'))
        self.assertTrue(ip4 is None)
        self.assertTrue(ip6 is None)


#!/usr/bin/python
# -*- coding: UTF-8 -*-

#import MySQLdb

import datetime
import pandas as pd
import seaborn as sns
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import random
import os
from matplotlib.pyplot import figure
from fastprogress import master_bar, progress_bar
import torch.nn as nn
import time
from sklearn.metrics import mean_squared_error
import warnings
import itertools
from kalmanfilter import KalmanFilter

sns.set_style("darkgrid")
warnings.filterwarnings("ignore")

class Predictor:
    def __init__(self, state_transition=1, process_noise=1, observation_model=1, observation_noise=5):
        # mysql db
        #self.db = MySQLdb.connect("localhost", "testuser", "test123", "TESTDB", charset='utf8' )

        self.kf = KalmanFilter(state_transition=state_transition, process_noise=process_noise,
           observation_model=observation_model, observation_noise=observation_noise)

    def preprocessing(self, ds):
        ds['进站时间'] = pd.to_datetime(ds['进站时间'])
        df = ds.set_index('进站时间')

        min_day = df.index.min().date().strftime('%Y-%m-%d')
        max_day = df.index.max().date().strftime('%Y-%m-%d')

        timl = pd.date_range(min_day, max_day, freq='1D')
        names = set(list(df['进站名称']))
        pairs = [i for i in itertools.combinations(names, 2)]
        result = {}
        print("processing data...")
        for dt in timl:
            record = dict.fromkeys(pairs , 0)
            fg = df[df.index.date == datetime.date(dt.year, dt.month, dt.day)]
            for n in names:
                x = fg[fg['进站名称'] == n].groupby('出站名称').count()['用户ID']
                if len(x) > 0:
                    for m,j in zip(x.to_frame().index, x):
                        if (m, n) in record:
                            record[(m,n)] += j
                        elif (n, m) in record:
                            record[(n,m)] += j
            result[dt] = record

        multiple_time_series = {}
        for k, v in result.items():
            for ke, ve in v.items():
                if ke not in multiple_time_series:
                    multiple_time_series[ke] = []
                multiple_time_series[ke].append(ve)
        return timl, multiple_time_series

    def do_predict(self, ob_pair, last_days=7, predict_days=3, sql_cmd="", mode="predict_all", output_dir="./output/"):
        """
        cursor = self.db.cursor()
        cursor.execute(sql_cmd) # SQL语句，用户自定义
        data = cursor.fetchone()
        # 做一些逻辑处理
        # ...
        # 数据格式pandas：id, 进站名称, 进站时间, 出战名称, 出站时间, 渠道编号, 价格
        """
        ds = pd.read_csv("trips.csv", encoding='gb18030')
        timl, multiple_time_series = self.preprocessing(ds)
        if ob_pair in multiple_time_series:
            pdata = multiple_time_series[ob_pair]
        else:
            ob_pair = (ob_pair[1], ob_pair[0])
            pdata = multiple_time_series[ob_pair]
        if mode == "predict_all":
            pred = []
            y = []
            index = timl[-(len(pdata)-last_days-1):]
            for i in range(len(pdata) - last_days - 1):
                x = self.kf.predict(pdata[:last_days + i], 1)
                r = x.observations.mean[0]
                pred.append(r)
                y.append(pdata[last_days + i + 1])
            plt.plot(index, y, 'b-', label="observation")
            plt.plot(index, pred, 'r-', label="predict")
            plt.title('Predict OD [{} and {}]'.format(ob_pair[0], ob_pair[1]),fontsize=21)
            plt.xticks(fontsize=18)
            plt.yticks(fontsize=18)
            plt.ylabel('Passageer Volume', fontsize=18)
            plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
            plt.gcf().autofmt_xdate()
            plt.legend()
            plt.savefig(output_dir + '/predict_image_all.jpg')
        elif mode == "predict_future":
            if predict_days > len(pdata):
                raise Exception
            new_stamps = []
            for i in range(1, predict_days+1):
                new_stamps.append(timl[-1] + datetime.timedelta(days=i))
            x = self.kf.predict(pdata, predict_days)
            r = 0.5 * np.array(x.observations.mean) + 0.5 * np.array(pdata[-3:][::-1])
            plt.plot(timl, pdata, 'r-', label="observation")
            plt.plot(new_stamps, r, 'b-', label="prediction")
            plt.title('Predict OD [{} and {}]'.format(ob_pair[0], ob_pair[1]), fontsize=21)
            plt.xticks(fontsize=18)
            plt.yticks(fontsize=18)
            plt.ylabel('Next {} Days Passageer Volume'.format(predict_days), fontsize=18)
            plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
            plt.gcf().autofmt_xdate()
            plt.legend()
            plt.savefig(output_dir + '/predict_image_future.jpg')
        else:
            raise Exception("Not support yet.")

if __name__ == "__main__":
    p_instance = Predictor()
    p_instance.do_predict(ob_pair=('Sta134', 'Sta63'))


# Mech Vib Tsak 2 (SDOF, Wind Vibration, Analytic Sol) #

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

    ## The required Libraries to be imported:

import numpy as np
from numpy import rad2deg
from numpy.linalg import inv
import matplotlib.pyplot as plt
from math import pi, sqrt, sin, cos, atan
from cmath import phase
from sympy import symbols
from sympy.physics.vector import init_vprinting, vlatex
init_vprinting(use_latex='mathjax', pretty_print=False)

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

    ## System Parameters:

        ### Materials:
E_s = 2.1e11 # Modulus of rigidity of steel [Pa]
E_c = 0.17e11 # Modulus of rigidity of concrete [Pa]
E = 0.2*E_s + 0.8*E_c # Total modulus of rigidity [Pa]
rho = 1000 # RC density [kg/m^3]

        ### Dimensions:
L = 5 # Floor length [m]
W = 5 # Floor width [m]
h = 0.5 # Floor thickness [m]
H  = 5 # Story hight [m]
l = 0.65 # Column length [m]
w = 0.65 # Column width [m]
V = L * W * h # Floor volume

        ### Dynamic Parametrs:
m = rho * V # Floor mass [kg]
I = (l * w**4) / 12 # Column inertia [m^4]
ko = (3 * E * I) / H**3 # Column stiffness (cantilever beam)
k = 4*ko # Equivelant columns stiffnesses
zeta = 0.05 # Daming ratio of the building
zeta_list = [0.05, 0.1, 0.15, 0.25, 0.5, 1, 2]
ze1, ze2, ze3, ze4, ze5, ze6, ze7 = symbols('zeta1, zeta2, zeta3, zeta4, zeta5, zeta6, zeta7')
ze_list = [ze1, ze2, ze3, ze4, ze5, ze6, ze7]
R = np.linspace(0.00001,4.00001,120) # Frequency ratio list
wn = sqrt(k/m) # Natural frequency
wd = sqrt(1-zeta**2)*wn # Damped Natural frequency
cc = 2*sqrt(m*k) # Critical damping
c = zeta * cc # Damping coefficient

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

    ## Input Data (Sin wave):

        ### Wind (f = Fo * cos(w_w*t), direct force):
freq_w = 10 # Excitation frequency [Hz]
w_w = 2 * pi * freq_w # Cicular frequency [rad/s]
v_w = 140 # Wind velocity [m/s]
rho_a = 1.2 # Air density [kg/m^3]
A_w = L*H # Wind faced area of the building
Fwo = 0.5 * rho_a * v_w**2 * A_w # Wind force [N]

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

    ## Simulation:

        ### Time Installation:
dt = 1/(8*freq_w) # Time step
end_time = 1 # Finish time
time = np.arange(0, end_time, dt) # Time range

        ### Initial Conditions:
xo = 0.0
vo = 0.0

        ### Response Histories & Amplitudes:
r_w = w_w/wn

Fw_history = [] # Wind input history

x_st = Fwo/k # Static defliction

Xw_history = [] # Wind response history
Xw = Fwo / ((k-m*w_w**2) + complex(0,w_w*c))

        ### Time Responses:
for t in time:
#--------------------------------------------------
    fw = Fwo * cos(w_w*t) # input wind force
    Fw_history.append(fw*1e-6) # MN
#--------------------------------------------------
    x_w = abs(Xw)*cos(w_w*t-phase(Xw)) # Steady state response from wind
    Xw_history.append(x_w*1e2) # cm
#------------------------------------------------------------------------------------------
M_history = np.zeros((len(zeta_list),len(R))) # Amplification Factor
PH_history = np.zeros((len(zeta_list),len(R))) # Phase
for i in range(0,len(zeta_list)):
    for j in range(0,len(R)):
        ze = zeta_list[i]
        r = R[j]
        M = 1 / sqrt( (1-r**2)**2 + (2*ze*r)**2 )
        M_history[i,j] = M
        th = rad2deg(atan(2*ze*r**2 / (1-r**2)))
        if th >= 0:
            PH = rad2deg(atan(2*ze*r**2 / (1-r**2)))
        elif th < 0:
            PH = 180+rad2deg(atan(2*ze*r**2 / (1-r**2)))
        else:
            PH = 0
        PH_history[i,j] = PH

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

    ## Results and Responses:

        ### Abs response x_eq(t):

fig1 = plt.figure()
rect = fig1.patch

graphW = fig1.add_subplot(1,1,1)
graphW.plot(time, Xw_history, linewidth=2.0)
graphW.plot(time, Fw_history, linewidth=1.0)
graphW.legend(['x_w(t) [cm]', 'f(t) [MN]'], fontsize=12, loc='lower right')
graphW.set_title(' "Wind System Response" ', fontsize=25)
graphW.set_xlabel('Time [s]', fontsize=20)
graphW.grid(True)

        ### Amplification Factor:

fig2a = plt.figure()
rect = fig2a.patch
fig2b = plt.figure()
rect = fig2b.patch

for i in range(0,len(zeta_list)):
    graphW2a = fig2a.add_subplot(1,1,1)
    graphW2a.plot(R, M_history[i,0:], linewidth=1.5)
    graphW2a.legend( ["${}$".format( vlatex(ze_list[k])+' = '+str(zeta_list[k])) for k in range(0,len(zeta_list))  ] , fontsize=13, loc='upper right' )
    graphW2a.set_title(' "Amplification factor with r for Wind" ', fontsize=25)
    graphW2a.set_ylabel('M = Xw/x_st', fontsize=20)
    graphW2a.set_xlabel('r = w_w/wn', fontsize=20)
    graphW2a.grid(True)
    graphW2b = fig2b.add_subplot(1,1,1)
    graphW2b.plot(R, PH_history[i,0:], linewidth=1.5)
    graphW2b.legend( ["${}$".format( vlatex(ze_list[k])+' = '+str(zeta_list[k])) for k in range(0,len(zeta_list))  ] , fontsize=13, loc='lower right' )
    graphW2b.set_title(' "Phase with r for Wind" ', fontsize=25)
    graphW2b.set_ylabel('Phase [deg]', fontsize=20)
    graphW2b.set_xlabel('r = w_w/wn', fontsize=20)
    graphW2b.grid(True)

plt.show()


#!/usr/bin/env python3
# json2parquet.py
# https://github.com/deadbits/vigil-llm
import os
import sys
import json
import argparse
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def load_json_file(file_path):
    try:
        with open(file_path, 'r') as f:
            data = json.load(f)
        logger.info(f"Successfully loaded JSON from {file_path}")
        return data

    except FileNotFoundError:
        logger.error(f"File {file_path} not found.")
        return None

    except json.JSONDecodeError:
        logger.error(f"Failed to decode JSON from {file_path}")
        return None

def convert_to_parquet(json_data):
    try:
        df = pd.DataFrame(json_data)
        table = pa.table(df)
        logger.info("Successfully converted JSON to Parquet")
        return table
    except Exception as e:
        logger.error(f"Conversion failed: {e}")
        return None

def save_parquet_to_disk(parquet_table, file_path):
    try:
        pq.write_table(parquet_table, file_path)
        logger.info(f"Successfully saved Parquet to {file_path}")
    except Exception as err:
        logger.error(f"Failed to save Parquet: {err}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '-f', '--file',
        help='json file',
        type=str,
        required=True
    )

    parser.add_argument(
        '-o', '--output',
        help='output file',
        type=str,
        required=True,
    )

    args = parser.parse_args()

    if not os.path.isfile(args.file):
        logger.error(f"File {args.file} not found.")
        sys.exit(1)

    if os.path.exists(args.output):
        logger.error(f"Ouput file {args.output} already exists.")
        sys.exit(1)

    json_data = load_json_file(args.file)

    if json_data:
        # Convert JSON to Parquet
        parquet_data = convert_to_parquet(json_data)

        if parquet_data:
            # Save Parquet back to disk
            save_parquet_to_disk(parquet_data, args.output)


import sys
sys.stdin = open("baek13702.txt")

"""
mid 만큼 나누어줄때
K 보다 적게 나오면
=> mid 를 줄인다.

K 보다 크거나 같게나오면 mid를 늘린다.

upper bound

최대 수치는 2**31 - 1 (약 20억)
"""

N, K = map(int, input().split())
arr = [0] * N
for i in range(0, N):
    arr[i] = int(input())

left = 1
right = 2**31

while left < right:
    mid = (left+right)//2

    cur = 0
    for i in range(0, N):
        # 남는건 버린다, 따라서 정수로 나누어떨어지는 만큼만 분배
        cur += (arr[i]//mid)

    if cur < K:
        right = mid
    else:
        left = mid+1
print(left-1)


import random
import matplotlib.pyplot as plt
import numpy as np

def TCPW(q, p, c, h, week):
    # initialization
    x_inv = q;  # inventory level
    x_dis = 0;  # remaining failure days
    T = 0;  # total cost
    PURCHASEPERMITTED = True;

    # update process
    for iter in range(week):
        if x_dis > 0:
            x_dis -= 1;
            PURCHASEPERMITTED = False;
        else:
            PURCHASEPERMITTED = True;

        if x_inv > 0:
            x_inv -= 1;
        else:
            T += p;

        if PURCHASEPERMITTED == True:
            u = random.uniform(0, 1);
            if u < 0.95:
                if x_inv <= q - 2:
                    x_inv += 2;
                    T += 2 * c;
                elif x_inv == q - 1:
                    x_inv += 1;
                    T += c;
            else:
                x_dis = 19;
        T += h * x_inv
    return T/week;
# factory produce equipment Model
# aim: determine the optimal inventory quantity threshold q

#-------------------------------------------------------------------------
#-------------------------------------------------------------------------

# total weeks
week = 1000000;
# threshold inventory quantity
Q = [i+1 for i in range(100)];
# parameter
p = 20;
c = 1;
h = 0.1;
# record for average cost
record_aver = [0] * len(Q);

count = 0;

for q in Q:
    aver_T = TCPW(q, p, c, h, week);
    record_aver[count] = aver_T;
    count += 1;

idx_opt = record_aver.index(min(record_aver));
quantity_opt = Q[idx_opt];
print('the optimal quantity is: ', quantity_opt);
print('the minimum average cost (per week) is: ', min(record_aver));

fig = plt.figure();
ax = fig.add_subplot(1,1,1);
ax.plot(Q, record_aver);
plt.title('Total Cost Per Week of each estimator q');
plt.xlabel('quantity');
plt.ylabel('Total Cost Per Week')
plt.show()


from nbformat import read
from rest_framework import serializers
from .models import Student

# Type3- validators
def start_with_r(value):
    if value[0].lower()!='r':
        raise serializers.ValidationError("Name Should start with R")

# #only this can handle serializers
# class StudentSerializer(serializers.ModelSerializer):
#     class Meta:
#         model = Student
#         fields = ['name','roll','city']

class StudentSerializer(serializers.ModelSerializer):
    # name = serializers.CharField(read_only=True) # to make this field read only -->method 1
    name = serializers.CharField(validators=[start_with_r]) # to assign a validator
    class Meta:
        model = Student
        fields = ['name','roll','city']
        #read_only_fields = ['name','roll'] # to make this field read only -->method 2
       # extra_kwargs = {'name':{'read_only':True}} # to make this field read only -->method 3

    # Type1 validation - field level validation
    # def validate_roll(self,value):
    #     if value>=200:
    #         raise serializers.ValidationError("Seats full")
    #     return value

    #Type2 - object validation
    def validate(self, data):
        name = data.get('name')
        roll = data.get('roll')
        city = data.get('city')
        if name.lower()=="dinesh" and city.lower()!='bhimsar':
            raise serializers.ValidationError("Dinesh is from bhimsar.pls fill it correctly")
        return data

    #     if name[0:2]!="19" or name[2:5]!="BCE" or len(name)!=9:
    #         raise serializers.ValidationError("name should be of type 19BCEXXXX")
    #     return data


import requests
import json
import re

# MITRE Groups https://attack.mitre.org/groups/
mitre_actors = requests.get('https://raw.githubusercontent.com/mitre/cti/master/enterprise-attack/enterprise-attack.json')
mitre_actors = mitre_actors.json()

mitre_actor_list = []
for adversary in mitre_actors['objects']:
    mitre_actor_dict = {}
    if adversary['type'] == 'intrusion-set':
        mitre_actor_dict['id'] = adversary['id']
        mitre_actor_dict['name'] = adversary['name']
        created = adversary['created']
        created = created.split('T')
        mitre_actor_dict['created'] = created[0]
        last_modified = adversary['modified']
        last_modified = last_modified.split('T')
        mitre_actor_dict['last_modified'] = last_modified[0]
        mitre_actor_dict['variations'] = []
        mitre_actor_dict['variations_custom'] = []
        external_references = adversary['external_references']
        for reference in external_references:
            if 'external_id' in reference.keys():
                mitre_actor_dict['url'] = reference['url']
        try:
            variations = adversary['aliases']
            for variation in variations:
                mitre_actor_dict['variations'].append(variation)
        except KeyError:
            mitre_actor_dict['variations'].append(adversary['name'])

        variations_custom = []
        for item in mitre_actor_dict['variations']:
            variation_upper = item.upper()
            variations_custom.append(variation_upper)
            if ' ' in variation_upper:
                variation_noSpace = variation_upper.replace(' ', '')
                variations_custom.append(variation_noSpace)
            if '-' in variation_upper:
                variation_noDash = variation_upper.replace('-', '')
                variations_custom.append(variation_noDash)
            try:
                if '-' in variation_noSpace:
                    variation_noSpace_noDash = variation_noSpace.replace('-', '')
                    variations_custom.append(variation_noSpace_noDash)
            except NameError:
                continue
        for i in variations_custom:
            mitre_actor_dict['variations_custom'].append(i)

        mitre_actor_list.append(mitre_actor_dict)

# Populate MITRE Group TTPs
for mitre_actor in mitre_actor_list:
    technique_list = []
    actorID = mitre_actor['id']
    ttp_patternIDs = []
    for object_relationships in mitre_actors['objects']:
        try:
            if object_relationships['source_ref'] == actorID:
                if 'attack-pattern' in object_relationships['target_ref']:
                    ttp_patternIDs.append(object_relationships['target_ref'])
        except KeyError:
            continue
    for patternID in ttp_patternIDs:
        for object_ttp in mitre_actors['objects']:
            if object_ttp['id'] == patternID:
                try:
                    for external_reference in object_ttp['external_references']:
                        try:
                            if 'CAPEC' not in external_reference['external_id']:
                                technique_list.append(external_reference['external_id'])
                        except KeyError:
                            continue
                except KeyError:
                    continue

    mitre_actor['TTPs'] = technique_list

# ETDA Actors https://apt.etda.or.th/cgi-bin/listgroups.cgi
etda_actors = requests.get('https://apt.etda.or.th/cgi-bin/getmisp.cgi?o=g')
etda_actors = etda_actors.json()

etda_actor_list = []
for adversary in etda_actors['values']:
    etda_actor_dict = {}
    etda_actor_dict['id'] = adversary['uuid']
    name = adversary['value']
    etda_actor_dict['name'] = name
    name = name.replace('[', '')
    name = name.replace(']', '')
    name_list = re.split(', |,', name)
    etda_actor_dict['variations'] = []
    etda_actor_dict['variations_custom'] = []
    metadata = adversary['meta']
    etda_actor_dict['url'] = 'https://apt.etda.or.th/cgi-bin/showcard.cgi?u=' + adversary['uuid']
    try:
        etda_actor_dict['created'] = metadata['date']
    except KeyError:
        etda_actor_dict['created'] = 'None Provided'
    try:
        for variation in metadata['synonyms']:
            variation = variation.replace('[', '')
            variation = variation.replace(']', '')
            etda_actor_dict['variations'].append(variation)
        for name_variation in name_list:
            if name_variation not in etda_actor_dict['variations']:
                etda_actor_dict['variations'].append(name_variation)
    except KeyError:
        for name_variation in name_list:
            if name_variation not in etda_actor_dict['variations']:
                etda_actor_dict['variations'].append(name_variation)
    if 'country' in metadata.keys():
        etda_actor_dict['country'] = metadata['country']
    if 'motivation' in metadata.keys():
        etda_actor_dict['motivation'] = metadata['motivation']
    if 'cfr-target-category' in metadata.keys():
        etda_actor_dict['targeted_industries'] = metadata['cfr-target-category']
    if 'cfr-suspected-victims' in metadata.keys():
        etda_actor_dict['targeted_countries'] = metadata['cfr-suspected-victims']

    variations_custom = []
    for item in etda_actor_dict['variations']:
        variation_upper = item.upper()
        variations_custom.append(variation_upper)
        if ' ' in variation_upper:
            variation_noSpace = variation_upper.replace(' ', '')
            variations_custom.append(variation_noSpace)
        if '-' in variation_upper:
            variation_noDash = variation_upper.replace('-', '')
            variations_custom.append(variation_noDash)
        try:
            if '-' in variation_noSpace:
                variation_noSpace_noDash = variation_noSpace.replace('-', '')
                variations_custom.append(variation_noSpace_noDash)
        except NameError:
            continue
    for i in variations_custom:
        etda_actor_dict['variations_custom'].append(i)

    etda_actor_list.append(etda_actor_dict)

# Comparison
merge_list = []
id_check = []
for mitre_actor in mitre_actor_list:
    for mitre_variation in mitre_actor['variations_custom']:
        for etda_actor in etda_actor_list:
            if mitre_variation in etda_actor['variations_custom']:
                # Start compiling final data
                merge_dict = {}
                if mitre_actor['id'] in id_check:
                    continue
                else:
                    id_check.append(mitre_actor['id'])

                    merge_dict['mitre_attack_id'] = mitre_actor['id']
                    merge_dict['mitre_attack_name'] = mitre_actor['name']
                    merge_dict['mitre_attack_aliases'] = mitre_actor['variations']
                    merge_dict['mitre_attack_created'] = mitre_actor['created']
                    merge_dict['mitre_attack_last_modified'] = mitre_actor['last_modified']
                    merge_dict['mitre_url'] = mitre_actor['url']
                    merge_dict['etda_id'] = etda_actor['id']
                    merge_dict['etda_name'] = etda_actor['name']
                    merge_dict['etda_aliases'] = etda_actor['variations']
                    merge_dict['etda_first_seen'] = etda_actor['created']
                    merge_dict['etda_url'] = etda_actor['url']
                    if 'country' in etda_actor.keys():
                        merge_dict['country'] = etda_actor['country']
                    else:
                        merge_dict['country'] = 'None Provided'
                    if 'motivation' in etda_actor.keys():
                        merge_dict['motivation'] = etda_actor['motivation']
                    else:
                        merge_dict['motivation'] = 'None Provided'
                    if 'targeted_industries' in etda_actor.keys():
                        merge_dict['victim_industries'] = etda_actor['targeted_industries']
                    else:
                        merge_dict['victim_industries'] = 'None Provided'
                    if 'targeted_countries' in etda_actor.keys():
                        merge_dict['victim_countries'] = etda_actor['targeted_countries']
                    else:
                        merge_dict['victim_countries'] = 'None Provided'
                    merge_dict['mitre_attack_ttps'] = mitre_actor['TTPs']

                    merge_list.append(merge_dict)

outfile = open('Categorized_Adversary_TTPs.json', 'w', encoding='utf-8')
json.dump(merge_list, outfile, indent=2, ensure_ascii=False)