#!/usr/bin/python3
# -*- coding: utf-8 -*-
# @Time    : 2022/6/6 17:07
# @Author  : 马赫
# @Email   : 1692303843@qq.com
# @FileName: olga_extract_tpl.py

import os
import shutil

try:
    from BiVFM.feature_engineering import feature_engineering
except Exception:
    pass

from BiVFM.olga_func import *
from BiVFM.config_transfer import *


def extract_tpl(autoolga_output_dir, tpl_dir):
    """
    extract the **`.tpl`** file from the project directory generated by AutoOLGAFleet
    :param autoolga_output_dir:olga project collection total path
    :param tpl_dir:`.tpl` file extraction path
    :return: None
    """

    try:
        shutil.rmtree(tpl_dir)
    except Exception:
        pass
    os.makedirs(tpl_dir)

    olga_proj_list = os.listdir(autoolga_output_dir)
    for olga_proj in olga_proj_list:
        origin_olga_tpl_filename = '%s/%s/%s.tpl' % (autoolga_output_dir, olga_proj, olga_proj)
        migrate_olga_tpl_filename = '%s/%s.tpl' % (tpl_dir, olga_proj)

        shutil.copy(origin_olga_tpl_filename, migrate_olga_tpl_filename)


def extract_csv(tpl_dir, raw_csv_dir):
    """
    extract the csv file from tpl list


    # 0	PT[PA](FLOWPATH_1.PIPE-1.1)"Pressure"
    # 1	TM[C](FLOWPATH_1.PIPE-1.1)"Fluid temperature"
    # 2	PT[PA](FLOWPATH_1.PIPE-1.3)"Pressure"
    # 3	TM[C](FLOWPATH_1.PIPE-1.3)"Fluid temperature"
    # 4	QG[M3/S](FLOWPATH_1.PIPE-1.6)"Gas volume flow"
    # 5	QLT[M3/S](FLOWPATH_1.PIPE-1.6)"Total liquid volume flow"
    # 6	QLTHL[M3/S](FLOWPATH_1.PIPE-1.6)"Volumetric flow rate oil"
    # 7	QLTWT[M3/S](FLOWPATH_1.PIPE-1.6)"Volumetric flow rate water"
    # 8	QT[M3/S](FLOWPATH_1.PIPE-1.6)"Total volume flow"
    # 9	GT[KG/S](PIPELINE.PIPE-1.2)"Total mass flow"
    # 10	PT[PA](PIPELINE.PIPE-1.2)"Pressure"
    # 11	QT[M3/S](PIPELINE.PIPE-1.2)"Total volume flow"
    # 12	TM[C](PIPELINE.PIPE-1.2)"Fluid temperature"
    # 13	GTSOUR[KG/S](SOURCE-1)"Source mass rate"

    :param tpl_dir:
    :param raw_csv_dir:
    :return:
    """

    try:
        shutil.rmtree(raw_csv_dir)
    except Exception:
        pass
    os.makedirs(raw_csv_dir)

    tpl_list2raw(tpl_dir=tpl_dir, raw_csv_dir=raw_csv_dir)

    _raw_data = pd.read_csv('%s/AUTO_0.csv' % raw_csv_dir, header=0,
                            index_col=0,
                            parse_dates=[0],  # The first column is the time series
                            date_parser=lambda x: df_parser(x, mode=0, unit='s', date_format='%Y-%m-%d %H:%M:%S')
                            )
    for num, i in enumerate(_raw_data.columns.to_list()):
        print('%s\t%s' % (num, i))

    print(_raw_data.shape)


def extract_xy(raw_csv_dir, xy_dir, pkl_dir, extract_xy, resample_time):
    """
    extract xy_dataset from raw_csv
    :param extract_xy: str 'dataset/feature_extractor.xlsx' base_config['extract_xy']
    :param raw_csv_dir:
    :param xy_dir:
    :param resample_time: str Time drop sampling, '2S'
    :return:
    """
    #
    # try:
    #     shutil.rmtree(xy_dir)
    # except Exception:
    #     pass
    # os.makedirs(xy_dir)
    # extract_x = pd.read_excel(extract_xy, sheet_name='x').iloc[:, :-1].dropna().values.T
    # extract_x[0] = extract_x[0].astype('int')
    # extract_x = [list(i) for i in extract_x]
    #
    # extract_y = pd.read_excel(extract_xy, sheet_name='y').iloc[:, :-1].dropna().values.T
    # extract_y[0] = extract_y[0].astype('int')
    # extract_y = [list(i) for i in extract_y]
    #
    # featrue_extract_dir(raw_csv_dir=raw_csv_dir, xy_csv_dir=xy_dir,
    #                     extract_x=extract_x, extract_y=extract_y)

    # read csv data
    csv_order = []
    raw_data_list = []

    for filename in os.listdir(xy_dir):
        if not filename.endswith('.csv'):
            continue
        csv_order.append(filename)

    # In ascending order between underscores and `.xlsx` letters
    csv_order = sorted(csv_order, key=lambda x: int(x[5: -4]))

    csv_order = csv_order[:]

    for filename in csv_order:
        path = os.path.join(xy_dir, filename)
        raw_data = read_df_second_resample(filename=path, resample_time=resample_time)

        # 特征工程变换
        try:

            raw_data_fe = feature_engineering(raw_data)
        except Exception as e:
            raw_data_fe = raw_data

        raw_data_list.append(raw_data_fe)

    # save raw_data_list as pkl file
    try:
        shutil.rmtree(pkl_dir)
    except Exception:
        pass
    os.makedirs(pkl_dir)
    pickle_save(path='%s/raw_data_list.pkl' % pkl_dir,
                content=raw_data_list)

    # raw_data_list = pickle_load('%s/raw_data_list.pkl' % pkl_dir)

    print('Number of sample files:', len(raw_data_list))


if __name__ == '__main__':

    step = 3
    base_yaml_path = 'olga_work1/conf_0_base.yaml'
    # base_yaml_path = r'D:\work\VFM\BiVFM\experiment\e2_multiwell\olga_work2\conf_0_base.yaml'

    base_conf = get_base_config(base_yaml_path)
    # print(base_conf)

    # Step 1: extract the **`.tpl`** file from the project directory generated by AutoOLGAFleet
    if step == 1:
        extract_tpl(base_conf['autoolga_output_dir'], base_conf['tpl_dir'])

    # Step 2: extract csv from tpl
    if step == 2:
        extract_csv(base_conf['tpl_dir'], base_conf['raw_csv_dir'])

    # Step 3: extract xy from csv and then form pkl
    if step == 3:
        extract_xy(base_conf['raw_csv_dir'], base_conf['xy_dir'], base_conf['pkl_dir'],
                   base_conf['extract_xy'],
                   base_conf['resample_time'])
