import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import scipy.interpolate as interp
import os

def ALGO_QSortPart(arr, start_idx, end_idx):
    i = start_idx - 1
    pivot = arr[end_idx]
    for j in range(start_idx, end_idx): 
        if arr[j] <= pivot: 
            i = i + 1
            arr[i], arr[j] = arr[j], arr[i]
    arr[i + 1], arr[end_idx] = arr[end_idx], arr[i + 1] 
    return (i + 1)
    
def ALGO_QSort(arr, start_idx, end_idx): 
    arr_temp = arr.copy()
    if start_idx < end_idx:
        pi = ALGO_QSortPart(arr_temp, start_idx, end_idx)
        arr_temp = ALGO_QSort(arr_temp, start_idx, pi - 1)
        arr_temp = ALGO_QSort(arr_temp, pi + 1, end_idx)
    return arr_temp

def genTempArr(start_temp = 0, end_temp = 50, step = 10, addition = np.array([25])):
    temp_arr = np.arange(start_temp, end_temp + step, step)
    temp_arr = np.concatenate((temp_arr, addition))
    temp_arr = ALGO_QSort(temp_arr, 0, len(temp_arr) - 1) * 1.0
    return temp_arr

def genSheetNameArr(start_temp = 0, end_temp = 50, step = 10, addition = np.array([25]), direct = "both"):
    assert(direct == "both" or direct == "fore" or direct == "back" or direct == "none"), \
            "TYPE should be \"both\", \"fore\", \"back\" or \"none\""
    temp_arr = genTempArr(start_temp, end_temp, step, addition)
    temp_arr_str = temp_arr.round(2).astype(str)
    temp_arr_str = np.char.replace(temp_arr_str, ".", "C")
    if direct == "none":
        sheet_name_arr = temp_arr_str
    elif direct == "fore" or direct == "back":
        sheet_name_arr = np.char.add(temp_arr_str, "_" + direct)
    else:
        sheet_name_arr_fore = np.char.add(temp_arr_str, "_+0+fore")
        sheet_name_arr_back = np.char.add(temp_arr_str, "_+1+back")
        sheet_name_arr = np.concatenate((sheet_name_arr_fore, sheet_name_arr_back))
        sheet_name_arr = ALGO_QSort(sheet_name_arr, 0, len(sheet_name_arr) - 1)
        sheet_name_arr = np.char.replace(sheet_name_arr, "+0+", "")
        sheet_name_arr = np.char.replace(sheet_name_arr, "+1+", "")
    return sheet_name_arr

def genExpDataExcel(file_name, sheet_name_arr, \
                    x_label = "Velocity", y_label = "DVout"):
    exp_blank_sheet = pd.DataFrame(columns = [x_label, y_label])
    with pd.ExcelWriter(file_name + ".xlsx") as writer:
        for sheet_name in sheet_name_arr:
            exp_blank_sheet.to_excel(writer, sheet_name, index = False)
    save_path = os.path.abspath(".") + "\\" + file_name + ".xlsx"
    print("Save to: " + save_path)
    return save_path
    