import unittest
import pandas as pd
import numpy as np
from helper_functions import DataFrameHelper, SeriesHelper

class HelperFunctionsTest(unittest.TestCase):
    """Unit Test Cases for helper_functions module"""

    def test_null_count(self):
        """Test that DataFrameHelper.null_count() returns total null cells"""
        df_test1 = pd.DataFrame([
                        [1, np.nan, 1],
                        [np.nan, 1, 1],
                        [1, 1, np.nan]
                       ])
        df_test2 = pd.DataFrame([
                        [1, 1, 1],
                        [1, 1, 1],
                        [1, 1, 1]
                       ])

        dfh_test1 = DataFrameHelper(df_test1)
        dfh_test2 = DataFrameHelper(df_test2)

        self.assertEqual(dfh_test1.null_count(), 3)
        self.assertEqual(dfh_test2.null_count(), 0)

    def test_train_test_split(self):
        """
        Test that DataFrameHelper.train_test_split() returns two dataframes
        with the appropriate fraction of the original dataframe
        """

        df_original = pd.DataFrame(
                                   np.random.randint(
                                                     0,
                                                     100,
                                                     size=(10, 4)),
                                   columns=list('ABCD'))

        frac = 0.8
        length_original = len(df_original)
        dfh_orginal = DataFrameHelper(df_original)

        df_train, df_test = dfh_orginal.train_test_split(frac)

        self.assertEqual((len(df_train) + len(df_test)), length_original)
        self.assertEqual(len(df_train), (length_original * frac))
        self.assertEqual(len(df_test), (length_original
                                        - (length_original * frac)))

    def test_randomize(self):
        """
        Test that the DataFrame.randomize() returns a dataframe not
        matching the original
        """

        df_original = pd.DataFrame(
                                   np.random.randint(
                                                     0,
                                                     100,
                                                     size=(10, 4)),
                                   columns=list('ABCD'))

        dfh_test1 = DataFrameHelper(df_original)
        dfh_test1.randomize()

        pd.testing.assert_frame_equal(df_original, dfh_test1.df)

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


## SETUP THE ENVIRONMENT

## Import general system modules
from google.colab import drive
from google_drive_downloader import GoogleDriveDownloader
from itertools import product
from itertools import chain
import random
import pickle
import gzip
import sys
import os

## Core modules
import pandas as pd
import numpy as np

## Import plotting libraries
from matplotlib import pyplot as plt
import matplotlib
import seaborn as sns

## ML modules
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA
from sklearn.cluster import DBSCAN
import umap

## UTILITY FUNCTIONS ##

## Define simple function that will truncate numeric variable within trunc_range
## This truncation is based on quantile function

def truncate_by_quantile(series, trunc_range = [0, 1]):
    wk_series = series.copy()
    q_values = series.quantile(trunc_range).values
    q_min = np.min(q_values)
    q_max = np.max(q_values)
    wk_series = np.where(np.logical_and(wk_series < q_min,~np.isnan(wk_series)) , q_min, wk_series)
    wk_series = np.where(np.logical_and(wk_series > q_max,~np.isnan(wk_series)) , q_max, wk_series)
    return(wk_series)

## Take single vm and plot it
def plot_single_vm(__perf_data, __vmid = None, transformation = None):

    ## Calculate dims
    __N = len(__perf_data)
    __VMs = list(__perf_data.keys())
    __T , __D = __perf_data[__VMs[0]].shape
    __F = list(__perf_data[__VMs[0]].columns)

    ##
    CHECK_VMS = any([i == __vmid for i in __VMs])

    ## Check vms
    if not CHECK_VMS:

        if __vmid is not None:
            print(__vmid, 'is not valid ID!\nRandom vm has been picked! ')

        vmIx = random.randint(0,(__N-1))
        __vmid = __VMs[vmIx]

    ###
    if transformation is None or transformation =='None':
        transformation  = return_transformer(None)

    ## Create instance of the data for this vm:
    __vmid, data  = sample_observation(__perf_data, vmid = __vmid)

    ## Create VMId + VM index title:
    plotTitle = 'VM ID:'+str(__vmid)

    ##  Create the plot:
    ts, axes =  plt.subplots(nrows=3, ncols=2, sharex=True, figsize=(30,8))
    ts.suptitle(plotTitle, fontsize=14)

    for i, ax in enumerate(axes.reshape(-1)):

        metric = __F[i]
        data_serices = data[metric]

        data_serices = transformation(data_serices)

        # For each metric find indexes with missing values
        missing_ix = data[data[metric].isna()].index
        minVal = data_serices.min()
        dummy_x = [minVal for i in range(len(missing_ix))]

        ax.set_ylabel(metric)
        ax.yaxis.set_label_position('right')
        ax.plot(data.index,data_serices)
        ax.scatter(missing_ix,dummy_x , c = 'red', s = 18, marker = 'x')
    plt.show()

## This is a utility function that samples the information for single virtual machine
## It can sample random as well as specific virtual machine and metric (metrics)
## It also can remove rows with missing values

def sample_observation(data, metrics = 'All',vmid = None):
    N = len(data)
    VMs = list(data.keys())
    if vmid == None:
        vmIx = random.randint(0,(N-1))
        vmid = VMs[vmIx]

    ## Selec the data
    data_sample = data[vmid].copy()

    if metrics != 'All':
        data_sample = data_sample.loc[: , data_sample.columns.intersection(metrics)]

    return(vmid, data_sample)

## These function setup some groundwork for easier training
## This function takes python dict where keys represent algorithm parameters
## the element of the list represent lists with parameter values
## eg for t-sne param_dict = {'perplexity':[5,10], 'n_components':[2]}
## the function creates generator with all the combinations of the parameter values

def model_params_product(pram_dict):
    keys = pram_dict.keys()
    vals = pram_dict.values()
    for instance in product(*vals):
        yield dict(zip(keys, instance))

## This is a function for model tunning
## it takes dataset, model function, hyperparameters, specific to this function
## and whether or not to apply PCA and how many pc to include
## other parameters are model parameters that are constant across all ins

def model_train(data, model , param_dict , apply_pca = False, n_pc = None, other_params = {}):

    ## Get the data shape
    Nrows, Ncols = data.shape
    model = str(model).upper()
    ## Init the arrays that we will use for modeling
    if isinstance(data, pd.DataFrame):
        data_instance = data.values
    else:
        data_instance = data

    n_features = [Ncols]

    ## Check if there should be pca
    ## and if all n_components are integer values

    if apply_pca:
        if n_pc == None:
            raise Exception('Provide int or list with number of components')

        if all([isinstance(x, int) for x in n_pc]):
            max_n_pc = max(n_pc)
            if max_n_pc > Ncols:
                print('The max number of PC({}) is greater than the number of features in the dataset ({}).\nThe max number of PC has been truncated to {}!\n'.format(max_n_pc,Ncols, Ncols))

            n_features = list(sorted(set([x if x <= Ncols else Ncols for x in n_pc ])))
            data_instance = PCA().fit_transform(data)

        else:
            raise Exception('Values', ', '.join([str(x) for x in  n_pc if not(isinstance(x, int))]), ' are not valid integers')

    models_ph = []

    i = 0

    for nfeat in n_features:
        ## Create generator with the hyperparams:
        model_params = model_params_product(param_dict)

        for model_param in model_params:
            params_for_tunning = model_param.copy()
            model_param.update(other_params)

            if model == 'DBSCAN':
                running_model = DBSCAN(**model_param).fit_predict(data_instance[:, :nfeat])
            elif model == 'TSNE':
                running_model = TSNE(**model_param).fit_transform(data_instance[:, :nfeat])
            elif model == 'UMAP':
                running_model = umap.UMAP(**model_param).fit_transform(data_instance[:, :nfeat])

            ## Add the number of PCA Components as additional argument
            if apply_pca:
                model_param['n_pc'] = nfeat
                params_for_tunning['n_pc'] = nfeat

            ## Log the progress
            print('Finished iteration ',i,' with parameters:', ', '.join([ k+':'+str(v) for k, v in model_param.items()]))
            i += 1

            ## Append the tuple of the parameters and the fitted model to list
            models_ph.append((params_for_tunning, running_model, other_params))

    print('Training has finished with {} tested combinations.'.format(len(models_ph)))

    return(models_ph)

## This is a plotting function that will help us visualize the results of our embedding.
## It takes a result tuple (as generated in the function model_train) and plots the first two dimensions of the embedding.

def plot_results(results, leading_param = None, color_array = None, force_categorical = False, point_size = 3):
    """
    This is simple plotting fuction that take tuning results object
    and yields plot grid. The grid can be colored by specific variable

   results: model results as generated from model_train function
   leading_param: the parameter that will serve for aligning the plots. leading_param is plotted row wise
   color_array: data, by which the plots will be colored. Can be single numpy vector or list of vectors with the same size to the results list
   force_categorical: Should coloring be represented on a categorical scale
    """

    ## Obtaining all distinct hyperparameters used for tuning:
    tuned_params = list(set(chain.from_iterable([list(r[0].keys()) for r in results])))

    ## Check if leading param is in the tuned_params
    lp_in_tuned_params = any([p  == leading_param for p in tuned_params])

    if leading_param == None or not(lp_in_tuned_params):
        leading_param = tuned_params[0]
        print('No valid leading parameter is given!')
        print('Valid tuned params are:', ', '.join(tuned_params))
        print(leading_param, 'is picked as leading param!')

    ## Check if the color variable
    if isinstance(color_array, list):
        if len(color_array) != 1 and len(color_array) != len(results):
            print(len(color_array))
            raise Exception('The length of the color array must be either 1 or equal to the length of the tuning results!')

    ## Obtain the distinct values of the leading param
    leading_param_values = [meta[leading_param] for meta, array, _ in results]
    leading_param_values_length = len(leading_param_values)
    ##
    leading_param_distinct_values = sorted(set(leading_param_values))
    leading_param_distinct_values_length = len(leading_param_distinct_values)

    ## Set the order of the subplots
    subplots_order = [results[k] for _, k in sorted(zip(leading_param_values, range(leading_param_values_length)))]
    orig_index = [k for _, k in sorted(zip(leading_param_values, range(leading_param_values_length)))]

    ## Define grid dims
    nrow = int(leading_param_distinct_values_length)
    ncol = int(leading_param_values_length/leading_param_distinct_values_length)

    ## If there is only 1 variable plot the results on a single row
    if ncol ==1:
        ncol = nrow
        nrow = 1

    ## Define plot:
    figsize_width =  (ncol  + 0.5 )* 6
    figsize_height = nrow * 6

    fig, ax = plt.subplots(nrows = nrow, ncols = ncol, figsize = (figsize_width,figsize_height))

    ## Init some variables
    color_var = None

    ## Plot
    for i, result in enumerate(subplots_order):
        meta_info = result[0]
        ylab = leading_param +':'+ str(meta_info[leading_param])
        col_ix = orig_index[i]
        title = ',   '.join([ k+':'+str(v) for k, v in meta_info.items() if k != leading_param])+'('+str(col_ix)+')'
        working_ax = ax.reshape(-1)[i]
        working_ax.set_title(title)
        working_ax.set_ylabel(ylab)

        ## Check if we should color
        ## Define the color variable

        if color_array != None:
            if len(color_array) == len(results):
                color_var  = color_array[col_ix]
            else:
                 color_var = color_array[0]

            ## Check if the color var is numeric or categorical:
            col_var_numeric = np.issubdtype(np.array(color_var).dtype, np.number)
            if col_var_numeric and not(force_categorical):
                myplt = working_ax.scatter(x = result[1][:,0],y = result[1][:,1], c = color_var,cmap = plt.cm.get_cmap('Spectral'), s = point_size)
                plt.colorbar(myplt, ax = working_ax)
            else:
                ## Define custom color palette:
                color_set_values = list(sorted(set(color_var)))
                color_set_length = len(color_set_values)
                color_set_values_indeces = list(range(color_set_length))

                color_set_index = dict([(str(v), i) for i, v in enumerate(color_set_values)])
                value_set_index = dict([(str(i), v) for i, v in enumerate(color_set_values)])
                color_var = [color_set_index[str(i)] for i in color_var]

                ## Apply label formatting:
                formatter = plt.FuncFormatter(lambda x ,loc: value_set_index[x])

                myplt = working_ax.scatter(x = result[1][:,0] ,y = result[1][:,1], c = color_var,cmap = plt.cm.get_cmap('Spectral', color_set_length), s = point_size)
                plt.colorbar(myplt, ax = working_ax,ticks = color_set_values_indeces) #, format = formatter
        else:
            myplt = working_ax.scatter(x = result[1][:,0] ,y = result[1][:,1], s = point_size)
    fig.show()

##########
def plot_embedding(result, fig_scale = 1, color_var_list = None, force_categorical = False, plot_centers = False, dont_plot_minus_one=True, point_size = 4):
    meta, emb, _ = result
    _, Ncomp = emb.shape

    plot_title = ', '.join([k+':'+str(v) for k, v in meta.items()])
    w_size, h_size =  (Ncomp-0.4)*6*fig_scale, (Ncomp-1)*6*fig_scale
    subtitle_y = round((Ncomp-1)/Ncomp,3) + (0.12/fig_scale)/Ncomp

    fig, ax = plt.subplots(nrows = Ncomp, ncols = Ncomp, figsize = (w_size, h_size), sharex=True, sharey=True)

    for i in range(Ncomp):
        for j in range(Ncomp):
            if i<=j:
                fig.delaxes(ax[i][j])
                continue

            working_ax = ax[i][j]

            if color_var_list is not None:
                color_var = color_var_list[0]
                col_var_numeric = np.issubdtype(np.array(color_var).dtype, np.number)

                if col_var_numeric and not(force_categorical):
                    myplt = working_ax.scatter(x = emb[:,j], y = emb[:,i], c = color_var,cmap = plt.cm.get_cmap('Spectral'), s = point_size)
                    plt.colorbar(myplt, ax = working_ax)

                else:
                    ## Define custom color palette:
                    color_set_values = list(sorted(set(color_var)))
                    color_set_length = len(color_set_values)
                    color_set_values_indeces = list(range(color_set_length))

                    color_set_index = dict([(str(v), i) for i, v in enumerate(color_set_values)])
                    value_set_index = dict([(str(i), v) for i, v in enumerate(color_set_values)])
                    color_var = [color_set_index[str(i)] for i in color_var]
                    myplt = working_ax.scatter(x = emb[:,j], y = emb[:,i], c = color_var,cmap = plt.cm.get_cmap('Spectral', color_set_length), s = point_size)

                    ## Apply label formatting:
                    formatter = plt.FuncFormatter(lambda x ,loc: value_set_index[str(x)])

                    plt.colorbar(myplt, ax = working_ax, ticks = color_set_values_indeces, format = formatter) #

                    ## Plot cluster centers:
                    if plot_centers:
                        ## Create ph for the calculated centers
                        calc_centers = []
                        for ctr in np.sort(np.unique(color_var)):
                            arr_ix = np.where(color_var == ctr)
                            str_label = value_set_index[str(ctr)]
                            if dont_plot_minus_one and str(str_label) == '-1':
                                continue

                            running_mean = np.mean(emb[arr_ix, :], axis  = 1, keepdims=False)[0]

                            calc_centers.append(running_mean)
                            working_ax.text(running_mean[j], running_mean[i], str_label
                                        ,   bbox=dict(boxstyle="square",
                                            ec=(1., 0.5, 0.5),
                                            fc=(1., 0.8, 0.8),
                                            ))

            else:
                working_ax.scatter(emb[:,j], emb[:,i], s = point_size)

            working_ax.set_xlabel('Component '+str(i+1))
            working_ax.set_ylabel('Component '+str(j+1))
    fig.tight_layout()
    #fig.suptitle()
    fig.suptitle(plot_title, y = subtitle_y, x = 0.05,horizontalalignment = 'left')
    plt.show()


import numpy as np       #the numpy library
import matplotlib.pyplot as plt     #tMatplotlib's pyplot

import sys    #gives access to a C-like sys lybrary
import os     #gives access to operating sys

print(sys.argv)   #prints any command line arguments, incl program name
print(os.getcwd()) #prints the current working directory


# To add a new cell, type '#%%'
# To add a new markdown cell, type '#%% [markdown]'
#%% Change working directory from the workspace root to the ipynb file location. Turn this addition off with the DataScience.changeDirOnImportExport setting
# ms-python.python added
import os
try:
	os.chdir(os.path.join(os.getcwd(), 'notebooks'))
	print(os.getcwd())
except:
	pass
#%%
from IPython import get_ipython

#%% [markdown]
#
#%% [markdown]
# Lambda School Data Science, Unit 2: Predictive Modeling
#
# # Applied Modeling, Module 2
#
# You will use your portfolio project dataset for all assignments this sprint.
#
# ## Assignment
#
# Complete these tasks for your project, and document your work.
#
# - [ ] Plot the distribution of your target.
#     - Classification problem: Are your classes imbalanced? Then, don't use just accuracy.
#     - Regression problem: Is your target skewed? If so, let's discuss in Slack.
# - [ ] Continue to clean and explore your data. Make exploratory visualizations.
# - [ ] Fit a model. Does it beat your baseline?
# - [ ] Try xgboost.
# - [ ] Get your model's permutation importances.
#
# You should try to complete an initial model today, because the rest of the week, we're making model interpretation visualizations.
#
#
# ## Reading
#
# Top recommendations in _**bold italic:**_
#
# #### Permutation Importances
# - _**[Kaggle / Dan Becker: Machine Learning Explainability](https://www.kaggle.com/dansbecker/permutation-importance)**_
# - [Christoph Molnar: Interpretable Machine Learning](https://christophm.github.io/interpretable-ml-book/feature-importance.html)
#
# #### (Default) Feature Importances
#   - [Ando Saabas: Selecting good features, Part 3, Random Forests](https://blog.datadive.net/selecting-good-features-part-iii-random-forests/)
#   - [Terence Parr, et al: Beware Default Random Forest Importances](https://explained.ai/rf-importance/index.html)
#
# #### Gradient Boosting
#   - [A Gentle Introduction to the Gradient Boosting Algorithm for Machine Learning](https://machinelearningmastery.com/gentle-introduction-gradient-boosting-algorithm-machine-learning/)
#   - _**[A Kaggle Master Explains Gradient Boosting](http://blog.kaggle.com/2017/01/23/a-kaggle-master-explains-gradient-boosting/)**_
#   - [_An Introduction to Statistical Learning_](http://www-bcf.usc.edu/~gareth/ISL/ISLR%20Seventh%20Printing.pdf) Chapter 8
#   - [Gradient Boosting Explained](http://arogozhnikov.github.io/2016/06/24/gradient_boosting_explained.html)
#   - _**[Boosting](https://www.youtube.com/watch?v=GM3CDQfQ4sw) (2.5 minute video)**_

#%%
import os, sys
in_colab = 'google.colab' in sys.modules

# If you're in Colab...
if in_colab:
    # Pull files from Github repo
    os.chdir('/content')
    get_ipython().system('git init .')
    get_ipython().system('git remote add origin https://github.com/LambdaSchool/DS-Unit-2-Applied-Modeling.git')
    get_ipython().system('git pull origin master')
    # Install packages in Colab
    get_ipython().system('pip install category_encoders==2.0.0')
    get_ipython().system('pip install eli5==0.10.1')
    get_ipython().system('pip install pandas-profiling==2.3.0')
    get_ipython().system('pip install pdpbox==0.2.0')
    get_ipython().system('pip install plotly==4.1.1')
    get_ipython().system('pip install shap==0.30.0')

    # Install required python packages
    get_ipython().system('pip install -r requirements.txt')

    # Change into directory for module
    os.chdir('module2')

#%%
# import important libraries
import numpy as np
import pandas as pd
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
import category_encoders as ce
from sklearn.ensemble import RandomForestClassifier
from sklearn.impute import SimpleImputer
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import confusion_matrix
from xgboost import XGBClassifier
from sklearn.metrics import roc_auc_score
import eli5
from eli5.sklearn import PermutationImportance

#%%
# import the dataset
df = pd.read_csv('https://raw.githubusercontent.com/oganm/dndstats/master/docs/uniqueTable.tsv',
                 sep='\t')

#%%
# double check the import worked correctly
df.head()

#%%
# create a has spells feature

# create a function that returns true if there is a spell or false if there is not
def check_spells(item):
  # check if the value is a float. This works because np.nan is a float value and the value will be a string if there are spells
  if isinstance(item, float):
    # if nan is found return false (no spells)
    return False
  else:
    # else return true (spell found)
    return True

# apply the function to the dataframe
df['has_spells'] = df['processedSpells'].apply(check_spells)

#%%
df.head()

#%%
# create a has feat feature

# create a function that returns true if it has a feat and false if there is not
def check_feats(item):
  if isinstance(item, float):
    return False
  else:
    return True

# apply the function to the dataframe
df['has_feats'] = df['feats'].apply(check_feats)

#%%
# create a HP per level feature
df['HP_per_level'] = df['HP'] / df['level']

#%%
df.head()

#%%
# drop uneeded columns due to high randomness or high variance
drop_columns_variance = ['name', 'date', 'day']
df = df.drop(columns = drop_columns_variance)

# drop columns due to leakage
drop_columns_leak = ['subclass', 'class']
df = df.drop(columns = drop_columns_leak)

# drop columns due to duplication
drop_columns_dup = ['race']
df = df.drop(columns = drop_columns_dup)

# check the head
df.head()

#%%
df['justClass'].value_counts()

#%%
# get our target
target = 'justClass'
targetAllowed = ['Fighter', 'Rogue', 'Cleric', 'Barbarian', 'Paladin', 'Ranger', 'Sorcerer', 'Wizard', 'Monk', 'Druid', 'Bard', 'Warlock']

#%%
# df = df[df[target] targetAllowed]
df = df.loc[df[target].isin(targetAllowed)]

#%%
df['justClass'].value_counts()

#%%
# get our features
train_features = df.drop(columns=target)
numeric_features = train_features.select_dtypes(include='number').columns.tolist()
cardinality = train_features.select_dtypes(exclude='number').nunique()
categorical_features = cardinality[cardinality <= 75].index.tolist()
features = numeric_features + categorical_features

#%%
features

#%%
# majority class check
majority_class = df[target].mode()
print('Majority class is', majority_class)
y_pred = [majority_class] * len(df)
accuracy_score(df[target], y_pred)

#%%
# train test split the data (80/20)
train, val = train_test_split(df, train_size=0.80, test_size=.20,
                               stratify=df[target], random_state=42)

#%%
X_train = train[features]
y_train = train[target]
X_val = val[features]
y_val = val[target]

#%%
# fit a pipeline (Decision Tree)
pipelineTree = make_pipeline(
    ce.OneHotEncoder(use_cat_names=True),
    SimpleImputer(strategy='mean'),
    StandardScaler(),
    DecisionTreeClassifier(max_depth=3)
)

pipelineTree.fit(X_train, y_train)

#%%
# validation accuracy (Decision Tree)
y_pred_tree = pipelineTree.predict(X_val)
print('Validation Accuracy', accuracy_score(y_val, y_pred_tree))

#%%
y_pred_tree

#%%
# fit a pipeline (Random Forest)
pipelineForest = make_pipeline(
    ce.OneHotEncoder(use_cat_names=True),
    SimpleImputer(strategy='mean'),
    StandardScaler(),
    RandomForestClassifier(n_estimators=50, random_state=42, n_jobs=-1)
)

pipelineForest.fit(X_train, y_train)

#%%
y_pred_forest = pipelineForest.predict(X_val)
print('Validation Accuracy', accuracy_score(y_val, y_pred_forest))

#%%
y_pred_forest

#%%
transformers = make_pipeline(
    ce.OneHotEncoder(use_cat_names=True),
    SimpleImputer(strategy='mean'),
    StandardScaler()
)

# tranform the data
X_train_transformed = transformers.fit_transform(X_train)
X_val_transformed = transformers.transform(X_val)

eval_set = [(X_train_transformed, y_train),
            (X_val_transformed, y_val)]

#%%
# model = RandomForestClassifier(n_estimators=50, random_state=42, n_jobs=-1)
# model.fit(X_train_transformed, y_train)

#%%

#%%
model = XGBClassifier(n_estimators=1000, n_jobs=-1)
model.fit(X_train_transformed, y_train, eval_set=eval_set, early_stopping_rounds=10)

#%%
# Validation accuracy
y_pred = model.predict(X_val_transformed)
print('Validation Accuracy', accuracy_score(y_val, y_pred))

#%%
y_pred

#%%
# permuter = PermutationImportance(
#     model,
#     scoring='accuracy',
#     n_iter=3,
#     random_state=42
# )

# permuter.fit(X_val_transformed, y_val)
# feature_names = X_val.columns.tolist()

# eli5.show_weights(
#     permuter,
#     top=None, # show importance of all features
#     feature_names=feature_names
# )

#%%
# the row to anaylze in the shaply plot
# change the number in row to find the shaply value
row = 25
row = X_train.iloc[[row]]

#%%
# import shap

# explainer = shap.TreeExplainer(model)
# row_processed = transformers.transform(row)
# shap_values = explainer.shap_values(row_processed)

# shap.initjs()
# shap.force_plot(
#     base_value=explainer.expected_value,
#     shap_values=shap_values,
#     features=row
# )

#%%

#%%


"""
測試自己的資料集，並存成檢測結果圖
"""

import torch, os, cv2
from model.model import parsingNet
from utils.common import merge_config
from utils.dist_utils import dist_print
import torch
import scipy.special, tqdm
import numpy as np
import torchvision.transforms as transforms
from data.dataset import LaneTestDataset
from data.constant import culane_row_anchor, tusimple_row_anchor

# 指定測試的配置訊息
backbone = "18" # 骨幹網路
dataset = "11" # 資料集類型
griding_num = 100 # 網格數
# test_model = "tusimple_18.pth" # 預訓練模型路徑 tusimple_18.pth
test_model = "放入欲訓練模型，自訓練or預訓練"
# data_root = "" # 開源資料集測試路徑
data_root = "自定義測試路徑" # 自定義測試路徑
data_save = "保存檢測結果" # 保存檢測結果

from PIL import Image
import os
import numpy as np
import glob

def loader_func(path):
    return Image.open(path)

class TestDataset(torch.utils.data.Dataset):
    def __init__(self, path, img_transform = None):
        super(TestDataset, self).__init__()
        self.path = path
        self.ing_transform = img_transform
        # self.list: 儲存測試圖片的相對路徑

    def __getitem__(self, index):
        name = glob.glob('%s/*.jpg'%self.path)[index]
        img = loader_func(name)

        if self.img_transform is not None:
            img = self.img_transform(img)
        return img, name

    def __len__(self):
        return len(self.list)

if __name__ == "__main__":
    torch.backends.cudnn.benchmark = True # 加速

    # args, cfg = merge_config() # 用終端機指定配置訊息
    dist_print('start testing...')
    assert backbone in ['18','34','50','101','152','50next','101next','50wide','101wide']

    if dataset == 'CULane':
        cls_num_per_lane = 18
    elif dataset == 'Tusimple':
        cls_num_per_lane = 56
    else:
    #     raise NotImplementedError
        cls_num_per_lane = 56

    net = parsingNet(pretrained = False, backbone=backbone, cls_dim = (griding_num+1, cls_num_per_lane, 4),
                    use_aux=False).cuda() # we dont need auxiliary segmentation in testing

    state_dict = torch.load(test_model, map_location='cpu')['model']
    compatible_state_dict = {}
    for k, v in state_dict.items():
        if 'module.' in k:
            compatible_state_dict[k[7:]] = v
        else:
            compatible_state_dict[k] = v

    net.load_state_dict(compatible_state_dict, strict=False)
    net.eval()
    # 圖像格式統一: (288, 800), 圖像張量, 歸一化(標準化?)
    img_transforms = transforms.Compose([
        transforms.Resize((288, 800)),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])
    if dataset == 'CULane':
        splits = ['test0_normal.txt', 'test1_crowd.txt', 'test2_hlight.txt', 'test3_shadow.txt', 'test4_noline.txt', 'test5_arrow.txt', 'test6_curve.txt', 'test7_cross.txt', 'test8_night.txt']
        datasets = [LaneTestDataset(data_root, os.path.join(data_root, 'list/test_split/'+split),img_transform = img_transforms) for split in splits]

        img_w, img_h = 1640, 590
        row_anchor = culane_row_anchor
    elif dataset == 'Tusimple':
        splits = ['test.txt']
        datasets = [LaneTestDataset(data_root,os.path.join(data_root, split),img_transform = img_transforms) for split in splits]
        img_w, img_h = 1280, 720
        row_anchor = tusimple_row_anchor
    else: # 自定義資料集
        # raise NotImplementedError
        datasets = TestDataset(data_root, img_transform = img_transforms)
        img_w, img_h = 1280, 720
        row_anchor = tusimple_row_anchor

    for dataset in zip(datasets): # split: 圖片列表 datasets: 統一格式之後的資料集
        loader = torch.utils.data.DataLoader(dataset, batch_size = 1, shuffle = False, num_workers = 1) # 載入資料集
        # fourcc = cv2.VideoWriter_fourcc(*"MJPG")
        # print(split[:-3] + "avi")
        # vout = cv2.VideoWriter(split[:-3] + "avi", fourcc, 30.0, (img_w, img_h)) # 保存結果成影片
        for i, data in enumerate(tqdm.tqdm(loader)): # 進度條顯示進度
            imgs, names = data # imgs: 圖像張量，圖像相對路徑
            imgs = imgs.cuda() # 使用GPU
            with torch.no_grad(): # 測試碼不計算梯度
                out = net(imgs) # 模型預測 輸出張量: [1,101,56,4]

            col_sample = np.linspace(0, 800 - 1, griding_num)
            col_sample_w = col_sample[1] - col_sample[0]

            out_j = out[0].data.cpu().numpy() # 數據類型轉換成numpy [101, 56, 4]
            out_j = out_j[:, ::-1, :] # 將第二維度倒著取[101, 56, 4]
            prob = scipy.special.softmax(out_j[:-1, :, :], axis=0) # [100, 56, 4]softmax計算(機率標準化到0-1之間且沿著維度0機率總和=1)
            idx = np.arange(griding_num) + 1 # 產生0-100
            idx = idx.reshape(-1, 1, 1) # [100, 1, 1]
            loc = np.sum(prob * idx, axis=0) # [56, 4]
            out_j = np.argmax(out_j, axis=0) # 返回最大值的索引
            loc[out_j == griding_num] = 0 # 若最大值的索引 = griding_num, 則歸零
            out_j = loc # [56, 4]

            # import pdb; pdb.set_trace()
            vis = cv2.imread(os.path.join(data_root,names[0])) # 讀取圖像[720, 1280, 3]
            for i in range(out_j.shape[1]): # 走遍所有列
                if np.sum(out_j[:, i] != 0) > 2: # 非0單位的格數<2
                    sum1 = np.sum(out_j[:, i] != 0)
                    for k in range(out_j.shape[0]): # 走遍所有行
                        if out_j[k, i] > 0:
                            ppp = (int(out_j[k, i] * col_sample_w * img_w / 800) - 1, int(img_h * (row_anchor[cls_num_per_lane-1-k]/288)) - 1 )
                            cv2.circle(vis,ppp,5,(0,255,0),-1)
            # 保存檢測結果圖
            cv2.imwrite(os.path.join(data_save, os.path.basename(names[0])), vis)

        # 保存影片結果
        #     vout.write(vis)

        # vout.release()


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

class grapher(object):
    """
    A simple class covering typical use of generating graphs.
    Exmple of usage:
    0) import class by using    "from matgrapher import grapher"
    1) create new object i.e.   "gr = grapher.grapher()"
    2) load labels              "gr.loadLabels(label1, label2)"
    3) load data                "gr.loadData(x_data1, y_data1, x_data2, y_data2)"
    4) generate graph           "gr.generateGraph()"
    5) remove loaded data       "gr.destroyGraphTable()"
    """

    def __init__(self):
        self.x_table = []
        self.y_table = []
        self.contour_plots = []
        self.point_table = [[], []]# x, y
        self.point_colors = [[], []]# color (in hex or matplotlib), alpha
        self.point_sizes = []
        self.point_alpha_change = []
        self.text_table = [[], []]
        self.labels = []
        self.xlim = []
        self.ylim = []
        self.linestyle = []
        self.colors = []
        self.show_label = []
        self.graphTitle = "Graph"
        self.axisNames = ["X Values", "Y Values"]
        self.outputFilename = "output/file.png"
        self.dpi = 300
        self.plotSize = [15*1.5/2.54, 15*1.5/2.54]
        self.showGrid = True
        self.saveFile = True
        self.showFigure = False
        self.logscale = 'none'

    def destroyGraphTable(self):
        '''
        Clean all data provided.
        '''
        for i in range(len(self.x_table)):
            del self.x_table[0]
        for i in range(len(self.y_table)):
            del self.y_table[0]
        for i in range(len(self.labels)):
            del self.labels[0]
        for i in range(len(self.linestyle)):
            del self.linestyle[0]
        for i in range(len(self.colors)):
            del self.colors[0]
        for i in range(len(self.point_table[0])):
            del self.point_table[0][0]
            del self.point_table[1][0]
        for i in range(len(self.point_colors[0])):
            del self.point_colors[0][0]
            del self.point_colors[1][0]
        for i in range(len(self.point_sizes)):
            del self.point_sizes[0]
        for i in range(len(self.point_alpha_change)):
            del self.point_alpha_change[0]
        for i in range(len(self.contour_plots)):
            del self.contour_plots[0]
        for i in range(len(self.text_table[0])):
            del self.text_table[0][0]
            del self.text_table[1][0]
        return None

    def loadLabels(self, label, *args):
        '''
        Load labels to internal array. Please provide them in order as x, y arguments were provided.
        Arguments:
        -> label (string) - label used in legend to describe a dataset,
        -> *args (string, ...) - following labels if needed to load more than one in one step
        '''
        if(len(self.labels)+len(args)+1<len(self.x_table)):
            warnings.warn(f"Not all data sets ({len(self.x_table)}) are labeled.")

        self.labels.append(label)
        self.show_label.append(True)
        if(len(args)>0):
            for i in range(len(args)):
                self.labels.append(args[i])
                self.show_label.append(True)

    def loadData(self, x_argument, y_argument, *args):
        '''
        Load data to internal tables. Please provide it in pairs [x1, y1, x2, y2, ...]
        Arguments:
        -> x_argument ([float, ...]) - one dimensional array of x axis values,
        -> y_argument ([float, ...]) - one dimensional array of y axis values.
        '''
        if(len(args)%2!=0):
            warnings.warn(f"Expected equal ammount of x and y data sets. Got ({int((len(args)+1)/2)}) and ({int((len(args)-1)/2)}).")

        self.x_table.append(x_argument.copy())
        self.y_table.append(y_argument.copy())
        if(len(args)>0):
            for i in range(int(len(args)/2)):
                self.x_table.append(args[2*i])
                self.y_table.append(args[2*i+1])

    def createContourPlot(self, fn, xlist, ylist):
        X, Y = np.meshgrid(xlist, ylist)
        Z = fn(X, Y)
        contour = [X, Y, Z]
        self.contour_plots.append(contour)

    def loadPoints(self, point, *args):
        '''
        Load points to internal table.
        If you wish to enable autocoloring of points, add "autocolor:[color of the point set],[opacity level]" string at the end of arguments.
        If opacity level is not provided, it will be assumed as 1.0.
        Arguments:
        -> point ([float, float]) - list containing point coordinates.
        '''
        autocolor_flag = False
        size_flag = False
        command_line = None #a command line at the end of arguments
        last_alpha = 0.0
        if(len(self.point_colors[1])>0):
            last_alpha = self.point_colors[1][-1]
        if(not isinstance(point[0], list) and not isinstance(point[0], np.ndarray)):
            self.point_table[0].append(point[0])
            self.point_table[1].append(point[1])
        else:
            if(len(point[0])!=len(point[1])):
                warnings.warn("Warning! Point data columns not equal in length! Not all points may be included.")
            for i in range(min([len(point[0]), len(point[1])])):
                self.point_table[0].append(point[0][i])
                self.point_table[1].append(point[1][i])
        if(len(args)>0):
            #check for command line at the end of arguments
            if(type(args[-1])==str):
                command_line = args[-1].split(";")
            #check for autocolor at the end of args
            for cmd in command_line:
                #if 'size' option is enabled
                if('size' in cmd):
                    size = cmd.split(":")[-1]
                    try:
                        size = float(size)
                    except:
                        size = 20.0# autosizing, if user fails to provide correct size
                    if(isinstance(point[0], list) or isinstance(point[0], np.ndarray)):
                        for i in range(min([len(point[0]), len(point[1])])):
                            self.point_sizes.append(size)
                    else:
                        self.point_sizes.append(size)
                    size_flag = True
                #if 'autocolor' option is enabled
                if('autocolor' in cmd):
                    color = cmd.split(":")[-1]#remove 'autocolor' from the argument
                    # fixing autocolor string
                    if(color == '' or color == 'autocolor'):# if user forgot to add color after autocolor
                        color = '#4e4e4e,1.0'
                    elif('autocolor' in color):#if 'autcolor' could not be removed (user forgot to add ':')
                        color = color.split("autocolor")[-1]# remove 'autocolor' forcefully
                    else:
                        # checking if user stated opacity correctly
                        try:
                            float(color.split(',')[-1])
                        except:
                            color = color.split(',')[0]+',1.0'
                    autocolor_flag = True
                    #colouring if the basic argument were list
                    if(isinstance(point[0], list) or isinstance(point[0], np.ndarray)):
                        for i in range(min([len(point[0]), len(point[1])])):
                            self.point_colors[0].append(color.split(',')[0])
                            self.point_colors[1].append(float(color.split(',')[1]))
                            if(last_alpha!=self.point_colors[1][-1]):
                                self.point_alpha_change.append(len(self.point_colors[1])-1)
                                last_alpha = self.point_colors[1][-1]
                    else:
                        self.point_colors[0].append(color.split(',')[0])
                        self.point_colors[1].append(float(color.split(',')[1]))
                        if(last_alpha!=self.point_colors[1][-1]):
                            self.point_alpha_change.append(len(self.point_colors[1])-1)
                            last_alpha = self.point_colors[1][-1]
            if(size_flag == False):
                if(isinstance(point[0], list) or isinstance(point[0], np.ndarray)):
                    for i in range(min([len(point[0]), len(point[1])])):
                        self.point_sizes.append(20.0)
                else:
                    self.point_sizes.append(20.0)
            #add additional points defined in the arguments
            for i in range(len(args)):
                if(autocolor_flag == True):
                    if(i == len(args)-1):
                        break
                    self.point_colors[0].append(color.split(',')[0])
                    self.point_colors[1].append(float(color.split(',')[1]))
                if(size_flag==True):
                    self.point_sizes.append(size)
                if(size_flag==False):
                    self.point_sizes.append(20.0)
                self.point_table[0].append(args[i][0])
                self.point_table[1].append(args[i][1])

    def changePointColor(self, color, point_index, end_point_index = None):
        '''
        Change color of a point or set of points in internal tables. If the point index is uknown, provide its position in form of a list.
        Arguments:
        -> color (string) - color of the point,
        -> point_index (int or [float, float]) - index number or position of the point,
        -> end_point_index (int or [float, float]) - index number or position of the last point of the set.
        '''
        if(isinstance(point_index, list)):
            pos = [var for var, val in enumerate(zip(self.point_table[0], self.point_table[1])) if val[0]==point_index[0] and val[1]==point_index[1]]
            for p in pos:
                self.point_colors[p] = color
            if(end_point_index != None):
                if(isinstance(end_point_index, list)):
                    end_pos = [var for var, val in enumerate(zip(self.point_table[0], self.point_table[1])) if val[0]==end_point_index[0] and val[1]==end_point_index[1]]
                    for i in range(min(pos), max(end_pos)):
                        self.point_colors[i] = color
                else:
                    for i in range(min(pos), end_point_index):
                        self.point_colors[i] = color
        else:
            self.point_colors[point_index] = color
            if(end_point_index != None):
                if(isinstance(end_point_index, list)):
                    end_pos = [var for var, val in enumerate(zip(self.point_table[0], self.point_table[1])) if val[0]==end_point_index[0] and val[1]==end_point_index[1]]
                    for i in range(point_index, max(end_pos)):
                        self.point_colors[i] = color
                else:
                    for i in range(point_index, end_point_index):
                        self.point_colors[i] = color

    def setPointColor(self, color, alpha = 1.0):
        '''
        Set color and transparency of a point.
        Arguments:
        -> color (string) - color of the point,
        -> alpha (float) - level of point transparency (between 0.0 and 1.0).
        '''
        self.point_colors[0].append(color)
        self.point_colors[1].append(alpha)

    def loadText(self, text, position, *args):
        '''
        Load positioned text into internal table. Please provide arguments in pairs of text1, position1, text2, position2, ...
        Arguments:
        -> text (string) - string to be displayed
        -> point ([float, float]) - list containing text position
        '''
        if(len(args)%2==1):
            warnings.warn(f"Warning! Data pairing (text, position) not complete. Expected even ammount of arguments, got odd. Last point will be set to [0.0, 0.0]")
        self.text_table[0].append(text)
        self.text_table[1].append(position)
        if(len(args)>0):
            for i in range(int(len(args)/2)):
                if(not isinstance(args[2*i+1], list)):
                    warnings.warn("Warning! Wrong data type provided! Expected position as list, got "+str(type(args[2*i+1]))+". Aborting loading text.")
                    return None
                self.text_table[0].append(args[2*i])
                self.text_table[1].append(args[2*i+1])
            if(i*2<len(args)-1 and i!=0):
                self.text_table[0].append(args[-1])
                self.text_table[1].append([0.0, 0.0])

    def hideLabel(self, label_index=False, label=''):
        '''
        Hide label from graph legend.
        Arguments:
        -> label_index (int or bool) - index of label to be hidden. Assign 'False' if more than one label is to be hidden.
        -> label (string) - label or labels to be hidden if share the same text.
        '''
        if(label!=''):
            if(label in self.labels):
                for i in range(len(self.labels)):
                    if(self.labels[i]==label):
                        self.show_label[i]=False
            else:
                warning.warn(f"Warning: No label {label} found!")
        else:
            if(label_index!=False):
                try:
                    self.show_label[label_index] = False
                except:
                    warning.warn(f"Warning: Could not hide label {label_index}. Size of label array is {len(self.show_label)}")

    def loadLineStyles(self, linestyle, *args):
        self.linestyle.append(linestyle)
        if(len(args)>0):
            for i in range(len(args)):
                self.linestyle.append(args[i])

    def changeLineStyle(self, index, newlinestyle, linestyle = '-'):
        if(index == 'u'):
            if(linestyle in self.linestyle):
                for i in range(len(self.linestyle)):
                    if(self.linestyle==linestyle):
                        self.linestyle = newlinestyle
        self.linestyle[index] = newlinestyle

    def loadColor(self, color, *args):
        self.colors.append(color)
        if(len(args)>0):
            for cl in args:
                self.colors.append(cl)

    def setAxisNames(self, X_axis, Y_axis):
        self.axisNames = [X_axis, Y_axis]

    def setGraphTitle(self, graph_title):
        self.graphTitle = graph_title

    def setFilename(self, filename):
        self.outputFilename = filename

    def setExportMethod(self, method):
        '''
        Sets method of exporting file.
        0 - save, don't show
        1 - don't save, show
        2 - save and show
        '''
        if(method!=0 and method!=1 and method!=2):
            warnings.warn(f"Warning: wrong export method provided: {method}. Falling back to method 1 (don't save, show)")
        if(method==0):
            self.saveFile = True
            self.showFigure = False
        if(method==1):
            self.saveFile = False
            self.showFigure = True
        if(method==2):
            self.saveFile = True
            self.showFigure = True

    def setGridVisibility(self, grid_visible):
        self.showGrid = grid_visible

    def setLogscaleMethod(self, logscale_method):
        self.logscale = logscale_method

    def generateGraph(self, data_x=None, data_y=None, axis_names=None, x_lim = None, y_lim = None, graph_title=None, line_styles=None, colors = None, legend=None, legend_args = '', legend_position = '', filename=None, dpi=None, plot_size = None , grid = None, save=None, show=None, tight_layout=True, log_scale = None):
        '''
        Draw a graph based on provided data.
        Arguments:
        -> data_x ([float, ...]) - data shown as argument X of the drawn chart,
        -> data_y ([float, ...]) - data shown as argument Y of the drawn chart,
        -> axis_names ([string, string]) - table of axis names ([0] - X axis, [1] - Y axis]),
        -> x_lim ([float, float]) - limits of x_axis (if an empty array is passed (default), no limits are imposed)
        -> y_lim ([float, float]) - limits of y_axis (if an empty array is passed (default), no limits are imposed)
        -> graph_title (string) - a title for drawn graph,
        -> line_styles ([string, ...]) - line styles as matplotlib argument
        -> colors ([string, ...]) - line colors as matplotlib argument
        -> legend ([string, ...]) - tabe of labels used in legend,
        -> legend_args (string) - arguments for plt.legend function
        -> filename (string) - name for the output file,
        -> dpi (int) - Dots Per Inch (resolution) of the exported graph,
        -> plot_size ([float, float]) - size of the exported graph (in inches - conversion ratio cm->inch is 1/2.54),
        -> grid (boolean) - argument for generating grid in the drawn graph
        -> save (boolean) - flag for saving drawn graph
        -> show (boolean) - flag for showing drawn graph
        '''
        #initializing graph arguments
        if(data_x == None):
            data_x = self.x_table
        if(data_y == None):
            data_y = self.y_table
        if(axis_names == None):
            axis_names = self.axisNames
        if(x_lim == None):
            x_lim = self.xlim
        if(y_lim == None):
            y_lim = self.ylim
        if(graph_title == None):
            graph_title = self.graphTitle
        if(line_styles == None):
            line_styles = self.linestyle
        if(colors == None):
            colors = self.colors
        if(legend == None):
            legend = self.labels
        if(filename == None):
            filename = self.outputFilename
        if(dpi == None):
            dpi = self.dpi
        if(plot_size == None):
            plot_size = self.plotSize
        if(grid == None):
            grid = self.showGrid
        if(save == None):
            save = self.saveFile
        if(show == None):
            show = self.showFigure
        if(log_scale == None):
            log_scale = self.logscale

        lx = len(data_x)
        ly = len(data_y)
        if(lx!=ly):
            raise Exception(f"Error: Expected equal length of data_x ({lx}) and data_y ({ly}) arrays.")
        if(lx>len(legend) and len(legend)!=0):
            warnings.warn("Warning: Not all provided data has been assigned with legend label.")
        if(lx<len(legend)):
            warnings.warn("Warning: Provided more labels than data.")

        plt.figure(figsize = (plot_size[0], plot_size[1]))
        for data_set_index, (xd, yd) in enumerate(zip(data_x, data_y)):
            if(data_set_index>=len(legend) or self.show_label[data_set_index]==False):
                if(data_set_index>=len(line_styles)):
                    if(data_set_index>=len(colors)):
                        plt.plot(xd, yd)
                    else:
                        plt.plot(xd, yd, colors[data_set_index])
                else:
                    if(data_set_index>=len(colors)):
                        plt.plot(xd, yd, linestyle=line_styles[data_set_index])
                    else:
                        plt.plot(xd, yd, linestyle=line_styles[data_set_index], color=colors[data_set_index])
            else:
                if(data_set_index>=len(line_styles)):
                    if(data_set_index>=len(colors)):
                        plt.plot(xd, yd, label = legend[data_set_index])
                    else:
                        plt.plot(xd, yd, label = legend[data_set_index], color=colors[data_set_index])
                else:
                    if(data_set_index>=len(colors)):
                        plt.plot(xd, yd, label = legend[data_set_index], linestyle=line_styles[data_set_index])
                    else:
                        plt.plot(xd, yd, label = legend[data_set_index], linestyle=line_styles[data_set_index], color=colors[data_set_index])
        if(len(legend)>0):
            if(legend_args==''):
                if(legend_position==''):
                    plt.legend()
                else:
                    plt.legend(loc=legend_position)
            else:
                if(legend_position==''):
                    plt.legend(legend_args)
                else:
                    plt.legend(legend_args, loc=legend_position)

        if(len(self.point_table[0])>0):
            if(len(self.point_colors[0])>0):
                if(len(self.point_sizes)>0):
                    for i in range(len(self.point_alpha_change)-1):
                        plt.scatter(self.point_table[0][self.point_alpha_change[i]:self.point_alpha_change[i+1]], self.point_table[1][self.point_alpha_change[i]:self.point_alpha_change[i+1]], c=self.point_colors[0][self.point_alpha_change[i]:self.point_alpha_change[i+1]], alpha = self.point_colors[1][self.point_alpha_change[i]], s=self.point_sizes[self.point_alpha_change[i]:self.point_alpha_change[i+1]])
                    plt.scatter(self.point_table[0][self.point_alpha_change[-1]:], self.point_table[1][self.point_alpha_change[-1]:], c=self.point_colors[0][self.point_alpha_change[-1]:], alpha = self.point_colors[1][self.point_alpha_change[-1]], s=self.point_sizes[self.point_alpha_change[-1]:])
                else:
                    for i in range(len(self.point_alpha_change)-1):
                        plt.scatter(self.point_table[0][self.point_alpha_change[i]:self.point_alpha_change[i+1]], self.point_table[1][self.point_alpha_change[i]:self.point_alpha_change[i+1]], c=self.point_colors[0][self.point_alpha_change[i]:self.point_alpha_change[i+1]], alpha = self.point_colors[1][self.point_alpha_change[i]])
                    plt.scatter(self.point_table[0][self.point_alpha_change[-1]:], self.point_table[1][self.point_alpha_change[-1]:], c=self.point_colors[0][self.point_alpha_change[-1]:], alpha = self.point_colors[1][self.point_alpha_change[-1]])
            else:
                plt.scatter(self.point_table[0], self.point_table[1])

        if(len(self.text_table)>0):
            for i in range(len(self.text_table[0])):
                plt.text(self.text_table[1][i][0], self.text_table[1][i][1], self.text_table[0][i])

        if(len(self.contour_plots)>0):
            for i in range(len(self.contour_plots)):
                cp = plt.contourf(self.contour_plots[i][0], self.contour_plots[i][1], self.contour_plots[i][2], cmap='coolwarm')
                plt.colorbar(cp)
                #plt.clabel(cp, inline=True)

        plt.xlabel(axis_names[0])
        plt.ylabel(axis_names[1])
        plt.title(graph_title)
        plt.grid(grid)
        if(len(x_lim)==2):
            plt.xlim(x_lim[0], x_lim[1])
        if(len(y_lim)==2):
            plt.ylim(y_lim[0], y_lim[1])
        if(log_scale=='y'):
            plt.yscale('log')
        if(log_scale=='x'):
            plt.xscale('log')
        if(log_scale=='xy'):
            plt.yscale('log')
            plt.xscale('log')
        if(save):
            if(tight_layout):
                plt.savefig(filename, bbox_inches='tight', dpi=dpi)
            else:
                plt.savefig(filename, dpi=dpi)
        if(show):
            plt.show()
        plt.close()


"""
============================================
The :mod:`mpi_array.globale_creation` Module
============================================

Defines :obj:`mpi_array.globale.gndarray` creation functions.

Ones and zeros
==============

.. autosummary::
   :toctree: generated/

   empty - Create uninitialised array.
   empty_like - Create uninitialised array same size/shape as another array.
   eye - Return 2D array with ones on diagonal and zeros elsewhere.
   identity - Return identity array.
   ones - Create one-initialised array.
   ones_like - Create one-initialised array same size/shape as another array.
   zeros - Create zero-initialised array.
   zeros_like - Create zero-initialised array same size/shape as another array.
   full - Create *fill value* initialised array.
   full_like - Create *fill value* initialised array same size/shape as another array.

From existing data
==================

.. autosummary::
   :toctree: generated/

   array - Returns :obj:`mpi_array.globale.gndarray` equivalent of input.
   asarray - Returns :obj:`mpi_array.globale.gndarray` equivalent of input.
   asanyarray - Returns :obj:`mpi_array.globale.gndarray` equivalent of input.
   copy - Create a replica of a specified array.

"""

from __future__ import absolute_import

import numpy as _np

from .license import license as _license, copyright as _copyright, version as _version
from . import logging as _logging  # noqa: E402,F401
from . import locale as _locale
from . import comms as _comms
from . import globale as _globale

__author__ = "Shane J. Latham"
__license__ = _license()
__copyright__ = _copyright()
__version__ = _version()

def empty(
    shape=None,
    dtype="float64",
    order='C',
    comms_and_distrib=None,
    intra_partition_dims=None,
    **kwargs
):
    """
    Creates array of uninitialised elements.

    :type shape: :samp:`None` or sequence of :obj:`int`
    :param shape: **Global** shape to be distributed amongst
       memory nodes.
    :type dtype: :obj:`numpy.dtype`
    :param dtype: Data type of array elements.
    :type order: :samp:`{'C', 'F'}`
    :param order: Only :samp:`'C'` implemented.
       Whether to store multi-dimensional data in row-major (C-style)
       or column-major (Fortran-style) order in memory.
    :type comms_and_distrib: :obj:`numpy.dtype`
    :param comms_and_distrib: Data type of array elements.
    :rtype: :obj:`mpi_array.globale.gndarray`
    :return: Newly created array with uninitialised elements.
    """
    if comms_and_distrib is None:
        comms_and_distrib = _comms.create_distribution(shape, **kwargs)
    lndarray_proxy, rma_window_buffer = \
        _locale.empty(
            comms_and_distrib=comms_and_distrib,
            dtype=dtype,
            order=order,
            return_rma_window_buffer=True,
            intra_partition_dims=intra_partition_dims
        )
    ary = \
        _globale.gndarray(
            comms_and_distrib=comms_and_distrib,
            rma_window_buffer=rma_window_buffer,
            lndarray_proxy=lndarray_proxy
        )

    return ary

def empty_like(ary, dtype=None, order='K', subok=True, **kwargs):
    """
    Return a new array with the same shape and type as a given array.

    :type ary: :obj:`numpy.ndarray`
    :param ary: Copy attributes from this array.
    :type dtype: :obj:`numpy.dtype`
    :param dtype: Specifies different dtype for the returned array.
    :type order: :samp:`{'C', 'F', 'A', or 'K'}`
    :param order: Only :samp:`'K'` implemented.
        Overrides the memory layout of the result. :samp:`'C'` means C-order, :samp:`'F'`
        means F-order, :samp:`'A'` means :samp:`'F'` if a is Fortran
        contiguous, :samp:`'C'` otherwise. :samp:`'K'` means match the layout
        of :samp:`{ary}` as closely as possible.
    :type subok: :obj:`bool`
    :param subok: Ignored.
       If True, then the newly created array will use the sub-class type of :samp:`{ary}`,
       otherwise it will be a base-class array. Defaults to True.
    :rtype: :samp:`type(ary)`
    :return: Array of uninitialized (arbitrary) data with the same shape and type as :samp:`{ary}`.
    """
    if dtype is None:
        dtype = ary.dtype
    if order == 'K':
        order = 'C'

    if (isinstance(ary, _globale.gndarray)):
        ret_ary = \
            empty(
                dtype=ary.dtype,
                comms_and_distrib=ary.comms_and_distrib,
                order=order,
                intra_partition_dims=ary.lndarray_proxy.intra_partition_dims
            )
    else:
        ary = _np.asanyarray(ary)
        ret_ary = empty(ary.shape, dtype=ary.dtype, order=order, **kwargs)

    return ret_ary

def full(
    shape=None,
    fill_value=0,
    *args,
    **kwargs
):
    """
    Return a new array of given shape and type, filled with :samp:`fill_value`.

    :type shape: :samp:`None` or sequence of :obj:`int`
    :param shape: **Global** shape to be distributed amongst
       memory nodes.
    :type fill_value: scalar
    :param fill_value: Fill value.
    :type dtype: :obj:`numpy.dtype`
    :param dtype: Data type of array elements.
    :type order: :samp:`{'C', 'F'}`
    :param order: Only :samp:`'C'` implemented.
       Whether to store multi-dimensional data in row-major (C-style)
       or column-major (Fortran-style) order in memory.
    :type comms_and_distrib: :obj:`numpy.dtype`
    :param comms_and_distrib: Data type of array elements.
    :rtype: :obj:`mpi_array.globale.gndarray`
    :return: Newly created array with uninitialised elements.
    """
    ary = empty(shape, *args, **kwargs)
    ary.fill_h(ary.dtype.type(fill_value))

    return ary

def full_like(ary, fill_value, *args, **kwargs):
    """
    Return a new array with the same shape and type as a given array.

    :type ary: :obj:`numpy.ndarray`
    :param ary: Copy attributes from this array.
    :type fill_value: scalar
    :param fill_value: Fill value.
    :type dtype: :obj:`numpy.dtype`
    :param dtype: Specifies different dtype for the returned array.
    :type order: :samp:`{'C', 'F', 'A', or 'K'}`
    :param order: Only :samp:`'K'` implemented.
        Overrides the memory layout of the result. :samp:`'C'` means C-order, :samp:`'F'`
        means F-order, :samp:`'A'` means :samp:`'F'` if a is Fortran
        contiguous, :samp:`'C'` otherwise. :samp:`'K'` means match the layout
        of :samp:`{ary}` as closely as possible.
    :type subok: :obj:`bool`
    :param subok: Ignored.
       If True, then the newly created array will use the sub-class type of :samp:`{ary}`,
       otherwise it will be a base-class array. Defaults to True.
    :rtype: :samp:`type(ary)`
    :return: Array of uninitialized (arbitrary) data with the same shape and type as :samp:`{ary}`.
    """
    ary = empty_like(ary, *args, **kwargs)
    ary.fill_h(ary.dtype.type(fill_value))

    return ary

def zeros(shape=None, dtype="float64", order='C', comms_and_distrib=None, **kwargs):
    """
    Creates array of zero-initialised elements.

    :type shape: :samp:`None` or sequence of :obj:`int`
    :param shape: **Global** shape to be distributed amongst
       memory nodes.
    :type dtype: :obj:`numpy.dtype`
    :param dtype: Data type of array elements.
    :type order: :samp:`{'C', 'F'}`
    :param order: Only :samp:`'C'` implemented.
       Whether to store multi-dimensional data in row-major (C-style)
       or column-major (Fortran-style) order in memory.
    :type comms_and_distrib: :obj:`numpy.dtype`
    :param comms_and_distrib: Data type of array elements.
    :rtype: :obj:`mpi_array.globale.gndarray`
    :return: Newly created array with zero-initialised elements.
    """
    return \
        full(
            shape=shape,
            fill_value=0,
            dtype=dtype,
            order=order,
            comms_and_distrib=comms_and_distrib,
            **kwargs
        )

def zeros_like(ary, *args, **kwargs):
    """
    Return a new zero-initialised array with the same shape and type as a given array.

    :type ary: :obj:`mpi_array.globale.gndarray`
    :param ary: Copy attributes from this array.
    :type dtype: :obj:`numpy.dtype`
    :param dtype: Specifies different dtype for the returned array.
    :type order: :samp:`{'C', 'F', 'A', or 'K'}`
    :param order: Only :samp:`'K'` implemented.
        Overrides the memory layout of the result. :samp:`'C'` means C-order, :samp:`'F'`
        means F-order, :samp:`'A'` means :samp:`'F'` if a is Fortran
        contiguous, :samp:`'C'` otherwise. :samp:`'K'` means match the layout
        of :samp:`{ary}` as closely as possible.
    :type subok: :obj:`bool`
    :param subok: Ignored.
       If True, then the newly created array will use the sub-class type of :samp:`{ary}`,
       otherwise it will be a base-class array. Defaults to True.
    :rtype: :obj:`mpi_array.globale.gndarray`
    :return: Array of zero-initialized data with the same shape and type as :samp:`{ary}`.
    """
    return full_like(ary, 0, *args, **kwargs)

def ones(shape=None, dtype="float64", comms_and_distrib=None, order='C', **kwargs):
    """
    Creates array of one-initialised elements.

    :type shape: :samp:`None` or sequence of :obj:`int`
    :param shape: **Global** shape to be distributed amongst
       memory nodes.
    :type dtype: :obj:`numpy.dtype`
    :param dtype: Data type of array elements.
    :type order: :samp:`{'C', 'F'}`
    :param order: Only :samp:`'C'` implemented.
       Whether to store multi-dimensional data in row-major (C-style)
       or column-major (Fortran-style) order in memory.
    :type comms_and_distrib: :obj:`numpy.dtype`
    :param comms_and_distrib: Data type of array elements.
    :rtype: :obj:`mpi_array.globale.gndarray`
    :return: Newly created array with one-initialised elements.
    """
    return \
        full(
            shape=shape,
            fill_value=1,
            dtype=dtype,
            order=order,
            comms_and_distrib=comms_and_distrib,
            **kwargs
        )

def ones_like(ary, *args, **kwargs):
    """
    Return a new one-initialised array with the same shape and type as a given array.

    :type ary: :obj:`mpi_array.globale.gndarray`
    :param ary: Copy attributes from this array.
    :type dtype: :obj:`numpy.dtype`
    :param dtype: Specifies different dtype for the returned array.
    :type order: :samp:`{'C', 'F', 'A', or 'K'}`
    :param order: Only :samp:`'K'` implemented.
        Overrides the memory layout of the result. :samp:`'C'` means C-order, :samp:`'F'`
        means F-order, :samp:`'A'` means :samp:`'F'` if a is Fortran
        contiguous, :samp:`'C'` otherwise. :samp:`'K'` means match the layout
        of :samp:`{ary}` as closely as possible.
    :type subok: :obj:`bool`
    :param subok: Ignored.
       If True, then the newly created array will use the sub-class type of :samp:`{ary}`,
       otherwise it will be a base-class array. Defaults to True.
    :rtype: :obj:`mpi_array.globale.gndarray`
    :return: Array of one-initialized data with the same shape and type as :samp:`{ary}`.
    """
    return full_like(ary, 1, *args, **kwargs)

def eye(N, M=None, k=0, dtype=_np.float):
    """
    Not implemented.
    Return a 2-D array with ones on the diagonal and zeros elsewhere.
    """
    raise NotImplementedError()

def identity(n, dtype=None):
    """
    Not implemented.
    Return the identity array.
    """
    raise NotImplementedError()

def copy(ary, **kwargs):
    """
    Return an array copy of the given object.

    :type ary: :obj:`mpi_array.globale.gndarray`
    :param ary: Array to copy.
    :type order: :samp:`{'C', 'F'}`
    :param order: Only :samp:`'C'` implemented.
       Whether to store multi-dimensional data in row-major (C-style)
       or column-major (Fortran-style) order in memory. Defaults to :samp:`'C'`.
    :rtype: :obj:`mpi_array.globale.gndarray`
    :return: A copy of :samp:`{ary}`.
    """
    return ary.copy(**kwargs)

def array(a, dtype=None, copy=True, order='K', subok=False, ndmin=0, **kwargs):
    """
    Create a :obj:`mpi_array.globale.gndarray` from an existing *array-like* object.

    :type object: array_like
    :param object: An array, any object exposing the array interface, an object
       whose :samp:`__array__` method returns an array, or any (nested) sequence.
    :type dtype: :obj:`numpy.dtype`
    :param dtype: The desired data-type for the array. If not given,
       then the type will be determined as the minimum type required to hold the
       objects in the sequence. This argument can only be used to `upcast` the
       array.
    :type copy: :obj:`bool`
    :param copy: If :samp:`True`, then the object is copied.
       Otherwise, a copy will only be made if :samp:`__array__` returns a copy,
       if :samp:`a` is a nested sequence, or if a copy is needed to satisfy any of
       the other requirements (:samp:`dtype`, :samp:`order`, etc.).
    :type order: :samp:`{'K', 'A', 'C', 'F'}`
    :param order: Only :samp:`C` implemented. Specify the memory layout of the array.
        If object is not an array, the newly created array will be in C order (row major)
    :type subok: :obj:`bool`
    :param subok: If :samp:`True`, then sub-classes will be passed-through, otherwise the
        returned array will be forced to be a base-class array.
    :type ndmin: int
    :param ndmin: Specifies the minimum number of dimensions that the resulting array should have.
        Ones will be pre-pended to the shape as needed to meet this requirement.
    :rtype: :obj:`mpi_array.globale.gndarray`
    :return: An array object satisfying the specified requirements.

    .. seealso:: :func:`asarray`, :func:`asanyarray`
    """
    if order == 'K':
        order = 'C'

    if hasattr(a, "__class__") and (a.__class__ is _globale.gndarray):
        if copy:
            ret_ary = a.copy()
        else:
            ret_ary = a
    elif isinstance(a, _globale.gndarray):
        if subok:
            ret_ary = a
        else:
            ret_ary =\
                _globale.gndarray(
                    comms_and_distrib=a.comms_and_distrib,
                    rma_window_buffer=a.rma_window_buffer,
                    lndarray_proxy=a.lndarray_proxy
                )
        if copy:
            ret_ary = a.copy()
    else:
        if "distrib_type" not in kwargs.keys() or kwargs["distrib_type"] is None:
            kwargs["distrib_type"] = _comms.DT_CLONED
        np_ary = _np.array(a, dtype=dtype, order=order, copy=False, subok=True, ndmin=ndmin)
        ret_ary = \
            empty(
                shape=np_ary.shape,
                dtype=np_ary.dtype,
                **kwargs
            )
        if (ret_ary.ndim == 0) and (ret_ary.locale_comms.have_valid_inter_locale_comm):
            ret_ary.lndarray_proxy.lndarray[...] = np_ary
        else:
            locale_rank_view_slice_n = ret_ary.lndarray_proxy.rank_view_slice_n
            if len(locale_rank_view_slice_n) > 0:
                globale_rank_view_slice_n = \
                    ret_ary.lndarray_proxy.locale_extent.locale_to_globale_slice_h(
                        locale_rank_view_slice_n
                    )
                ret_ary.lndarray_proxy.lndarray[locale_rank_view_slice_n] =\
                    np_ary[globale_rank_view_slice_n]

        ret_ary.intra_locale_barrier()

    if ret_ary.ndim < ndmin:
        ret_ary = ret_ary.reshape((1,) * (ndmin - ret_ary.ndim) + tuple(ret_ary.shape))

    return ret_ary

def asarray(a, dtype=None, order=None, **kwargs):
    """
    Converts :samp:`{a}` (potentially via a copy)
    to a :obj:`mpi_array.globale.gndarray`.
    The :samp:`{kwargs}` are as for the :func:`mpi_array.comms.create_distributon` function
    and determine the distribution for the
    returned :obj:`mpi_array.globale.gndarray`.

    :type a: scalar, :obj:`tuple`, :obj:`list`, :obj:`numpy.ndarray`, etc
    :param a: Object converted to a :obj:`mpi_array.globale.gndarray`.
    :type dtype: :obj:`numpy.dtype`
    :param dtype: The :obj:`numpy.dtype` for the returned :obj:`mpi_array.globale.gndarray`.
    :type order: :samp:`{'C', 'F'}`
    :param order: Only :samp:`'C'` implemented.
       Whether to store multi-dimensional data in row-major (C-style)
       or column-major (Fortran-style) order in memory. Defaults to :samp:`'C'`.
    :rtype: :obj:`mpi_array.globale.gndarray`
    :return: The object :obj:`a` converted to an instance
       of :obj:`mpi_array.globale.gndarray`.

    .. seealso:: :func:`array`, :func:`asanyarray`
    """
    return array(a, dtype, copy=False, order=order, **kwargs)

def asanyarray(a, dtype=None, order=None, **kwargs):
    """
    Convert the input to an ndarray, but pass :obj:`mpi_array.globale.gndarray` subclasses through.

    :type a: scalar, :obj:`tuple`, :obj:`list`, :obj:`numpy.ndarray`, etc
    :param a: Object converted to a :obj:`mpi_array.globale.gndarray`.
    :type dtype: :obj:`numpy.dtype`
    :param dtype: The :obj:`numpy.dtype` for the returned :obj:`mpi_array.globale.gndarray`.
    :type order: :samp:`{'C', 'F'}`
    :param order: Only :samp:`'C'` implemented.
       Whether to store multi-dimensional data in row-major (C-style)
       or column-major (Fortran-style) order in memory. Defaults to :samp:`'C'`.
    :rtype: :obj:`mpi_array.globale.gndarray`
    :return: The object :obj:`a` converted to an instance
       of :obj:`mpi_array.globale.gndarray`.

    .. seealso:: :func:`array`, :func:`asarray`
    """
    return array(a, dtype, copy=False, order=order, subok=True, **kwargs)

__all__ = [s for s in dir() if not s.startswith('_')]


import pytest

import pandas as pd
import numpy  as np
import iirsBenchmark.regressors as regressors

from sklearn.base           import RegressorMixin
from sklearn.exceptions     import NotFittedError
from scipy.optimize         import check_grad
from sklearn.utils._testing import ignore_warnings
from sklearn.exceptions     import ConvergenceWarning

ds_names = pd.read_csv(
    './datasets/FeynmanEquations.csv')['Filename'].values

@pytest.mark.parametrize("regressor", regressors.__all__)
def test_has_properties(regressor):

    # checking attributes does not need to instantiate the class
    regressor_class = getattr(regressors, regressor)

    assert hasattr(regressor_class, 'stochastic_executions')
    assert hasattr(regressor_class, 'interpretability_spectrum')
    assert hasattr(regressor_class, 'grid_params')

@pytest.mark.parametrize("regressor", regressors.__all__)
def test_is_regressorMixin(regressor):

    regressor_class = getattr(regressors, regressor)
    regressor_instance = regressor_class()

    assert isinstance(regressor_instance, RegressorMixin)

@ignore_warnings(category=ConvergenceWarning)
@pytest.mark.parametrize("regressor,ds_name", zip(
    regressors.__all__,
    np.random.choice(ds_names, len(regressors.__all__))
))
def test_fit_and_predict(regressor, ds_name):

    # testing with random datasets. The dataset correctness should be
    # verified in test_feynman.

    data = pd.read_csv(
            f'./datasets/train/{ds_name}_UNI.csv', sep=',',
            header=0, index_col=False).values

    X, y = data[:, :-1], data[:, -1]

    regressor_class = getattr(regressors, regressor)
    # creating with small configurations (only the regressors with
    # slow fitting process)
    predictor = regressor_class(**{
        # ITEA
        'gens' : 10,
        'popsize' : 10,

        # stop warnings when label is missing
        'labels' : [f'x_{i}' for i in range(X.shape[1])],

        # MLP
        'hidden_layer_sizes' : (50,),
        'activation' : 'identity',
    })

    # Should fail, not fitted yet
    with pytest.raises(NotFittedError):
        predictor.predict(X)

    # should succeed
    predictor.fit(X, y)

    predictions = predictor.predict(X)

    assert predictions.shape == (len(y), )
    assert np.any(np.isfinite(predictions))

@pytest.mark.parametrize("regressor,ds_name", zip(
    ['ITEA_regressor', 'Linear_regressor', 'Lasso_regressor'],
    ds_names[:3]
))
def test_gradients(regressor, ds_name):

    data = pd.read_csv(
        f'./datasets/train/{ds_name}_UNI.csv', sep=',',
        header=0, index_col=False).values

    X, y = data[:, :-1], data[:, -1]

    # not using random datasets because ITEA may need a proper configuration
    # to find a solution that does not present discontinuity. The first 3
    # are simple enough
    regressor_class = getattr(regressors, regressor)

    predictor = regressor_class(**{
        # ITEA
        'gens' : 25,
        'popsize' : 25,

        # stop warnings when label is missing
        'labels' : [f'x_{i}' for i in range(X.shape[1])],
    })
    predictor.fit(X, y)

    # auxiliary functions: check_grad takes a single parameter function
    pred_aux = lambda x: predictor.predict(np.array(x).reshape(1, -1))[0]
    grad_aux = lambda x: predictor.gradients(np.array(x).reshape(1, -1))[0]

    for x in X:
        assert check_grad(pred_aux, grad_aux, x, epsilon=1e-4) < np.std(y)

@pytest.mark.parametrize("regressor", regressors.__all__)
def test_to_string(regressor):

    regressor_class = getattr(regressors, regressor)
    regressor_instance = regressor_class()

    # methods used in the experiments. All regressors should have at least
    # those methods
    for method in ['fit', 'predict', 'to_str', 'score']:
        # should throw exception if the class does not have to_string
        hasattr(regressor_instance, method)

        assert callable(getattr(regressor_instance, method))


# -*- coding: utf-8 -*-
"""
Various small and named graphs, together with some compact generators.

"""
__author__ ="""Aric Hagberg (hagberg@lanl.gov)\nPieter Swart (swart@lanl.gov)"""
#    Copyright (C) 2004-2008 by
#    Aric Hagberg <hagberg@lanl.gov>
#    Dan Schult <dschult@colgate.edu>
#    Pieter Swart <swart@lanl.gov>
#    All rights reserved.
#    BSD license.

__all__ = ['make_small_graph',
           'LCF_graph',
           'bull_graph',
           'chvatal_graph',
           'cubical_graph',
           'desargues_graph',
           'diamond_graph',
           'dodecahedral_graph',
           'frucht_graph',
           'heawood_graph',
           'house_graph',
           'house_x_graph',
           'icosahedral_graph',
           'krackhardt_kite_graph',
           'moebius_kantor_graph',
           'octahedral_graph',
           'pappus_graph',
           'petersen_graph',
           'sedgewick_maze_graph',
           'tetrahedral_graph',
           'truncated_cube_graph',
           'truncated_tetrahedron_graph',
           'tutte_graph']

import networkx as nx
from networkx.generators.classic import empty_graph, cycle_graph, path_graph, complete_graph
from networkx.exception import NetworkXError

#------------------------------------------------------------------------------
#   Tools for creating small graphs
#------------------------------------------------------------------------------
def make_small_undirected_graph(graph_description, create_using=None):
    """
    Return a small undirected graph described by graph_description.

    See make_small_graph.
    """
    if create_using is not None and create_using.is_directed():
        raise NetworkXError("Directed Graph not supported")
    return make_small_graph(graph_description, create_using)

def make_small_graph(graph_description, create_using=None):
    """
    Return the small graph described by graph_description.

    graph_description is a list of the form [ltype,name,n,xlist]

    Here ltype is one of "adjacencylist" or "edgelist",
    name is the name of the graph and n the number of nodes.
    This constructs a graph of n nodes with integer labels 0,..,n-1.

    If ltype="adjacencylist"  then xlist is an adjacency list
    with exactly n entries, in with the j'th entry (which can be empty)
    specifies the nodes connected to vertex j.
    e.g. the "square" graph C_4 can be obtained by

    >>> G=nx.make_small_graph(["adjacencylist","C_4",4,[[2,4],[1,3],[2,4],[1,3]]])

    or, since we do not need to add edges twice,

    >>> G=nx.make_small_graph(["adjacencylist","C_4",4,[[2,4],[3],[4],[]]])

    If ltype="edgelist" then xlist is an edge list
    written as [[v1,w2],[v2,w2],...,[vk,wk]],
    where vj and wj integers in the range 1,..,n
    e.g. the "square" graph C_4 can be obtained by

    >>> G=nx.make_small_graph(["edgelist","C_4",4,[[1,2],[3,4],[2,3],[4,1]]])

    Use the create_using argument to choose the graph class/type.
    """
    ltype=graph_description[0]
    name=graph_description[1]
    n=graph_description[2]

    G=empty_graph(n, create_using)
    nodes=G.nodes()

    if ltype=="adjacencylist":
        adjlist=graph_description[3]
        if len(adjlist) != n:
            raise NetworkXError("invalid graph_description")
        G.add_edges_from([(u-1,v) for v in nodes for u in adjlist[v]])
    elif ltype=="edgelist":
        edgelist=graph_description[3]
        for e in edgelist:
            v1=e[0]-1
            v2=e[1]-1
            if v1<0 or v1>n-1 or v2<0 or v2>n-1:
                raise NetworkXError("invalid graph_description")
            else:
                G.add_edge(v1,v2)
    G.name=name
    return G

def LCF_graph(n,shift_list,repeats,create_using=None):
    """
    Return the cubic graph specified in LCF notation.

    LCF notation (LCF=Lederberg-Coxeter-Fruchte) is a compressed
    notation used in the generation of various cubic Hamiltonian
    graphs of high symmetry. See, for example, dodecahedral_graph,
    desargues_graph, heawood_graph and pappus_graph below.

    n (number of nodes)
      The starting graph is the n-cycle with nodes 0,...,n-1.
      (The null graph is returned if n < 0.)

    shift_list = [s1,s2,..,sk], a list of integer shifts mod n,

    repeats
      integer specifying the number of times that shifts in shift_list
      are successively applied to each v_current in the n-cycle
      to generate an edge between v_current and v_current+shift mod n.

    For v1 cycling through the n-cycle a total of k*repeats
    with shift cycling through shiftlist repeats times connect
    v1 with v1+shift mod n

    The utility graph K_{3,3}

    >>> G=nx.LCF_graph(6,[3,-3],3)

    The Heawood graph

    >>> G=nx.LCF_graph(14,[5,-5],7)

    See http://mathworld.wolfram.com/LCFNotation.html for a description
    and references.

    """
    if create_using is not None and create_using.is_directed():
        raise NetworkXError("Directed Graph not supported")

    if n <= 0:
        return empty_graph(0, create_using)

    # start with the n-cycle
    G=cycle_graph(n, create_using)
    G.name="LCF_graph"
    nodes=G.nodes()

    n_extra_edges=repeats*len(shift_list)
    # edges are added n_extra_edges times
    # (not all of these need be new)
    if n_extra_edges < 1:
        return G

    for i in range(n_extra_edges):
        shift=shift_list[i%len(shift_list)] #cycle through shift_list
        v1=nodes[i%n]                    # cycle repeatedly through nodes
        v2=nodes[(i + shift)%n]
        G.add_edge(v1, v2)
    return G

#-------------------------------------------------------------------------------
#   Various small and named graphs
#-------------------------------------------------------------------------------

def bull_graph(create_using=None):
    """Return the Bull graph. """
    description=[
        "adjacencylist",
        "Bull Graph",
        5,
        [[2,3],[1,3,4],[1,2,5],[2],[3]]
        ]
    G=make_small_undirected_graph(description, create_using)
    return G

def chvatal_graph(create_using=None):
    """Return the Chvátal graph."""
    description=[
        "adjacencylist",
        "Chvatal Graph",
        12,
        [[2,5,7,10],[3,6,8],[4,7,9],[5,8,10],
         [6,9],[11,12],[11,12],[9,12],
         [11],[11,12],[],[]]
        ]
    G=make_small_undirected_graph(description, create_using)
    return G

def cubical_graph(create_using=None):
    """Return the 3-regular Platonic Cubical graph."""
    description=[
        "adjacencylist",
        "Platonic Cubical Graph",
        8,
        [[2,4,5],[1,3,8],[2,4,7],[1,3,6],
         [1,6,8],[4,5,7],[3,6,8],[2,5,7]]
        ]
    G=make_small_undirected_graph(description, create_using)
    return G

def desargues_graph(create_using=None):
    """ Return the Desargues graph."""
    G=LCF_graph(20, [5,-5,9,-9], 5, create_using)
    G.name="Desargues Graph"
    return G

def diamond_graph(create_using=None):
    """Return the Diamond graph. """
    description=[
        "adjacencylist",
        "Diamond Graph",
        4,
        [[2,3],[1,3,4],[1,2,4],[2,3]]
        ]
    G=make_small_undirected_graph(description, create_using)
    return G

def dodecahedral_graph(create_using=None):
    """ Return the Platonic Dodecahedral graph. """
    G=LCF_graph(20, [10,7,4,-4,-7,10,-4,7,-7,4], 2, create_using)
    G.name="Dodecahedral Graph"
    return G

def frucht_graph(create_using=None):
    """Return the Frucht Graph.

    The Frucht Graph is the smallest cubical graph whose
    automorphism group consists only of the identity element.

    """
    G=cycle_graph(7, create_using)
    G.add_edges_from([[0,7],[1,7],[2,8],[3,9],[4,9],[5,10],[6,10],
                [7,11],[8,11],[8,9],[10,11]])

    G.name="Frucht Graph"
    return G

def heawood_graph(create_using=None):
    """ Return the Heawood graph, a (3,6) cage. """
    G=LCF_graph(14, [5,-5], 7, create_using)
    G.name="Heawood Graph"
    return G

def house_graph(create_using=None):
    """Return the House graph (square with triangle on top)."""
    description=[
        "adjacencylist",
        "House Graph",
        5,
        [[2,3],[1,4],[1,4,5],[2,3,5],[3,4]]
        ]
    G=make_small_undirected_graph(description, create_using)
    return G

def house_x_graph(create_using=None):
    """Return the House graph with a cross inside the house square."""
    description=[
        "adjacencylist",
        "House-with-X-inside Graph",
        5,
        [[2,3,4],[1,3,4],[1,2,4,5],[1,2,3,5],[3,4]]
        ]
    G=make_small_undirected_graph(description, create_using)
    return G

def icosahedral_graph(create_using=None):
    """Return the Platonic Icosahedral graph."""
    description=[
        "adjacencylist",
        "Platonic Icosahedral Graph",
        12,
        [[2,6,8,9,12],[3,6,7,9],[4,7,9,10],[5,7,10,11],
         [6,7,11,12],[7,12],[],[9,10,11,12],
         [10],[11],[12],[]]
        ]
    G=make_small_undirected_graph(description, create_using)
    return G

def krackhardt_kite_graph(create_using=None):
    """
    Return the Krackhardt Kite Social Network.

    A 10 actor social network introduced by David Krackhardt
    to illustrate: degree, betweenness, centrality, closeness, etc.
    The traditional labeling is:
    Andre=1, Beverley=2, Carol=3, Diane=4,
    Ed=5, Fernando=6, Garth=7, Heather=8, Ike=9, Jane=10.

    """
    description=[
        "adjacencylist",
        "Krackhardt Kite Social Network",
        10,
        [[2,3,4,6],[1,4,5,7],[1,4,6],[1,2,3,5,6,7],[2,4,7],
         [1,3,4,7,8],[2,4,5,6,8],[6,7,9],[8,10],[9]]
         ]
    G=make_small_undirected_graph(description, create_using)
    return G

def moebius_kantor_graph(create_using=None):
    """Return the Moebius-Kantor graph."""
    G=LCF_graph(16, [5,-5], 8, create_using)
    G.name="Moebius-Kantor Graph"
    return G

def octahedral_graph(create_using=None):
    """Return the Platonic Octahedral graph."""
    description=[
        "adjacencylist",
        "Platonic Octahedral Graph",
        6,
        [[2,3,4,5],[3,4,6],[5,6],[5,6],[6],[]]
        ]
    G=make_small_undirected_graph(description, create_using)
    return G

def pappus_graph():
    """ Return the Pappus graph."""
    G=LCF_graph(18,[5,7,-7,7,-7,-5],3)
    G.name="Pappus Graph"
    return G

def petersen_graph(create_using=None):
    """Return the Petersen graph."""
    description=[
        "adjacencylist",
        "Petersen Graph",
        10,
        [[2,5,6],[1,3,7],[2,4,8],[3,5,9],[4,1,10],[1,8,9],[2,9,10],
         [3,6,10],[4,6,7],[5,7,8]]
        ]
    G=make_small_undirected_graph(description, create_using)
    return G

def sedgewick_maze_graph(create_using=None):
    """
    Return a small maze with a cycle.

    This is the maze used in Sedgewick,3rd Edition, Part 5, Graph
    Algorithms, Chapter 18, e.g. Figure 18.2 and following.
    Nodes are numbered 0,..,7
    """
    G=empty_graph(0, create_using)
    G.add_nodes_from(range(8))
    G.add_edges_from([[0,2],[0,7],[0,5]])
    G.add_edges_from([[1,7],[2,6]])
    G.add_edges_from([[3,4],[3,5]])
    G.add_edges_from([[4,5],[4,7],[4,6]])
    G.name="Sedgewick Maze"
    return G

def tetrahedral_graph(create_using=None):
    """ Return the 3-regular Platonic Tetrahedral graph."""
    G=complete_graph(4, create_using)
    G.name="Platonic Tetrahedral graph"
    return G

def truncated_cube_graph(create_using=None):
    """Return the skeleton of the truncated cube."""
    description=[
        "adjacencylist",
        "Truncated Cube Graph",
        24,
        [[2,3,5],[12,15],[4,5],[7,9],
         [6],[17,19],[8,9],[11,13],
         [10],[18,21],[12,13],[15],
         [14],[22,23],[16],[20,24],
         [18,19],[21],[20],[24],
         [22],[23],[24],[]]
        ]
    G=make_small_undirected_graph(description, create_using)
    return G

def truncated_tetrahedron_graph(create_using=None):
    """Return the skeleton of the truncated Platonic tetrahedron."""
    G=path_graph(12, create_using)
#    G.add_edges_from([(1,3),(1,10),(2,7),(4,12),(5,12),(6,8),(9,11)])
    G.add_edges_from([(0,2),(0,9),(1,6),(3,11),(4,11),(5,7),(8,10)])
    G.name="Truncated Tetrahedron Graph"
    return G

def tutte_graph(create_using=None):
    """Return the Tutte graph."""
    description=[
        "adjacencylist",
        "Tutte's Graph",
        46,
        [[2,3,4],[5,27],[11,12],[19,20],[6,34],
         [7,30],[8,28],[9,15],[10,39],[11,38],
         [40],[13,40],[14,36],[15,16],[35],
         [17,23],[18,45],[19,44],[46],[21,46],
         [22,42],[23,24],[41],[25,28],[26,33],
         [27,32],[34],[29],[30,33],[31],
         [32,34],[33],[],[],[36,39],
         [37],[38,40],[39],[],[],
         [42,45],[43],[44,46],[45],[],[]]
        ]
    G=make_small_undirected_graph(description, create_using)
    return G


import nose.tools as nt       # Using nosetests
import numpy as np
from pendulum import Pendulum

def test_special_method_call_in_Pendulum_class():
    """Tests that theta and omega is computed correctly."""
    # Test values
    theta    = np.pi/4 # Angular position of the pendulum
    omega    = 0.1     # Angular velocity of the pendulum
    analytic = [0.1, -3.1530534197454685]
    eps      = 10**(-7)

    pendel   = Pendulum(L=2.2)
    computed = pendel(0,[theta,omega])

    assert(abs(computed[0] - analytic[0]) < eps)
    assert(abs(computed[1] - analytic[1]) < eps)

def test_special_method_call_in_Pendulum_class_keeps_a_peldelum_at_rest():
    """Tests that the pendulum is kept at rest."""
    # Test values
    theta0   = 0
    omega0   = 0
    analytic = [0, 0]
    eps      = 10**(-7)

    pendel   = Pendulum()
    computed = pendel(0,[theta0,omega0])

    assert(abs(computed[0] - analytic[0]) < eps)
    assert(abs(computed[1] - analytic[1]) < eps)

@nt.raises(AttributeError)
def test_error_if_solve_method_has_not_been_called():
    """
    Test that the solve method has been called. Error raised if attributes dont exist.
    """
    pendel = Pendulum()
    theta  = pendel.theta
    omega  = pendel.omega
    time   = pendel.t

def test_only_the_latest_solution_is_stored():
    """Tests that latest solution overwrites previous ones."""
    y0_1 = [0, 0]
    T_1  = 5
    dt_1 = 0.1

    y0_2 = [2, 3]
    T_2  = 15
    dt_2 = 0.01

    y0_3 = [1, 4]
    T_3  = 10
    dt_3 = 0.05

    pendel = Pendulum()
    pendel.solve(y0_1, T_1, dt_1)
    len_1 = len(pendel.t) #store previous length
    pendel.solve(y0_2, T_2, dt_2)
    len_2 = len(pendel.t) #store previous length
    pendel.solve(y0_3, T_3, dt_3)
    #Check length of t
    assert(len(pendel.t) != len_1)
    assert(len(pendel.t) != len_2)

    pendel2 = Pendulum()
    pendel2.solve(y0_3, T_3, dt_3)
    # Solve pendel2 for case #3 only
    # Check so that pendel is the latest solution
    for i in range(len(pendel.x)):
        assert(pendel.x[i] == pendel2.x[i])
        assert(pendel.y[i] == pendel2.y[i])

def test_solve_method_in_Pendulum_class_theta_omega_zero_arrays():
    """
    Test solve method keeps pendulum at rest for initial y0=[0,0] while t=i*dt.
    """
    y0 = [0, 0]
    T  = 5
    dt = 0.1

    pendel = Pendulum()
    pendel.solve(y0, T, dt)

    for i in range(len(pendel.t)):
        assert(pendel.t[i]     == i*pendel.dt)
        assert(pendel.theta[i] == 0)
        assert(pendel.omega[i] == 0)

def test_x_and_y_positions_are_correct():
    """
    Tests that x and y position is computed correctly by testing x^2 + y^2 = L^2.
    """
    y0  = [2, 3]
    T   = 15
    dt  = 0.01
    eps = 10**(-7)

    pendel = Pendulum(L=2)
    pendel.solve(y0, T, dt)
    sol = pendel.solution

    array_of_L = np.zeros(len(sol.y[0])) + (pendel.L**2)
    computed_radius_squared = pendel.x**2 + pendel.y**2
    for i in range(len(sol.y[0])):
        assert(abs(computed_radius_squared[i] - array_of_L[i]) < eps)

if __name__ == "__main__":
    import nose
    nose.run()


from matplotlib.pyplot import figure, plot, savefig, title, xlabel, ylabel
from numpy import genfromtxt, linspace
from sys import argv

input_file = str(argv[1])
output_file = str(argv[2])

data = genfromtxt(input_file, delimiter=',')
x = data[0]
sol = data[1]

fig = figure()
title('Solution')
plot(x, sol, '-')
xlabel('x')
ylabel('sol')
savefig(output_file)

print('Generated ' + output_file)


"""Import statements tell python what packages you'll be using.

You can use 'as' to change how you refer to the package. In this file,
I import matplotlib.pyplot as plt so I don't have to type out
'matplotlib.pyplot' every time I want to use it.
"""
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

###############################################################################
"""After the import statements, I define function(s).
Function definitions start with the word 'def' and end when
the indentation ends. For example:

def ex1():
    print("This line is part of the function 'ex1' because it's indented")
    print("This line is still part of the function 'ex1'")
    print("Blank lines do not end the current level of indentation")

print("This line is not part of the function 'ex1' because it's not indented")

"""

###############################################################################
# These next lines 'call' the functions defined above.
# If you erase the lines below (or put a # in front of them) this code won't
# 'do' anything (it won't generate a graph anymore.)