# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
# default_exp tabular.preprocessing
# -

#hide
from nbdev.showdoc import *

# # Preprocessor for tabular models

#export
from fastai.tabular.all import *
from fastai.data.all import *
from fastai.vision.data import get_grid
from enveco.data.las import *
from enveco.data.image import *
import matplotlib.patches as mpl_patches
from typing import Tuple
from fastai.metrics import *

# ## Preprocessing

# Preprocess csv-files of our field plot data into `TabularPandas` to feed into models.

# +
# export

from tqdm import tqdm
tqdm.pandas()

class EnvecoPreprocessor():
    "Needs a bit refactoring"
    def __init__(self, train_path, valid_path, test_path, **kwargs):
        self.train_df = pd.read_csv(train_path)
        self.train_df = self.train_df.rename(columns = lambda x: re.sub('[\.]+', '_', x))
        self.valid_df = pd.read_csv(valid_path)
        self.valid_df = self.valid_df.rename(columns = lambda x: re.sub('[\.]+', '_', x))
        self.test_df = pd.read_csv(test_path)
        self.test_df = self.test_df.rename(columns = lambda x: re.sub('[\.]+', '_', x))
        self.train_df['is_valid'] = 0
        self.valid_df['is_valid'] = 1
        self.train_val_df = pd.concat((self.train_df, self.valid_df))
    
    def preprocess_lidar(self, target_col, path, min_h:float=1.5, mask_plot:bool=True, normalize:bool=True,
                         log_y:bool=False, save_path:str=None) -> Tuple[TabularPandas, TabularPandas]:
        "Preprocess data and return (train_val, test) -tuple. Optionally log-transform target column with np.log1p"
        trainval = self.train_val_df.copy()
        test = self.test_df.copy()
        feature_cols = point_cloud_metric_cols
        trainval[point_cloud_metric_cols] = trainval.progress_apply(lambda row: point_cloud_metrics(f'{path}/{row.sampleplotid}.las',
                                                                                                    row.x, row.y, 
                                                                                                    min_h=min_h, 
                                                                                                    mask_plot=mask_plot), 
                                                                    axis=1, result_type='expand')
        test[point_cloud_metric_cols] = test.progress_apply(lambda row: point_cloud_metrics(f'{path}/{row.sampleplotid}.las',
                                                                                            row.x, row.y, 
                                                                                            min_h=min_h, 
                                                                                            mask_plot=mask_plot), 
                                                            axis=1, result_type='expand')
        
        if log_y:
            trainval[target_col] = np.log1p(trainval[target_col])
            test[target_col] = np.log1p(test[target_col])
            
        procs = None
        if normalize:
            procs = [Normalize]#.from_stats(*norm_stats)]
        trainval_tb = TabularPandas(trainval, procs=procs,
                                    cont_names=feature_cols, y_names=target_col,
                                    splits=ColSplitter(col='is_valid')(trainval))
        test_tb = TabularPandas(test, procs=procs,
                                cont_names=feature_cols, y_names=target_col)
        if save_path:
            trainval.to_csv(f'{save_path}/las_trainval.csv', index=False)
            test.to_csv(f'{save_path}/las_test.csv', index=False)
            with open(f'{save_path}/las_features.txt', 'w') as f:
                f.writelines("%s\n" % c for c in feature_cols)
        return trainval_tb, test_tb
    
    def load_las(self, path, target_col, normalize:bool=True, log_y:bool=False) -> Tuple[TabularPandas, TabularPandas]:
        "Load previously preprocessed las data"
        trainval = pd.read_csv(f'{path}/las_trainval.csv')
        test = pd.read_csv(f'{path}/las_test.csv')
        with open(f'{path}/las_features.txt', 'r') as f:
            feature_cols = [c.rstrip() for c in f.readlines()]
    
        if log_y:
            trainval[target_col] = np.log1p(trainval[target_col])
            test[target_col] = np.log1p(test[target_col])
        procs = None
        if normalize:
            procs = [Normalize]#.from_stats(*norm_stats)]
        trainval_tb = TabularPandas(trainval, procs=procs,
                                    cont_names=feature_cols, y_names=target_col,
                                    splits=ColSplitter(col='is_valid')(trainval))
        test_tb = TabularPandas(test, procs=procs,
                                cont_names=feature_cols, y_names=target_col)
        return trainval_tb, test_tb
    
    def preprocess_image(self, target_col, path, radius:int=31, mask_plot:bool=True, normalize:bool=True,
                         log_y:bool=False, save_path:str=None) -> Tuple[TabularPandas, TabularPandas]:
        "Preprocess dataframes and return (train_val, test) -tuple"
        trainval = self.train_val_df.copy()
        test = self.test_df.copy()
        #feature_cols = image_metric_cols
        trainval_feats = []
        for s in tqdm(trainval.sampleplotid.unique()):
            feats = process_image_features(f'{path}/{s}.tif', mask_plot, radius)
            feats['sampleplotid'] = s
            trainval_feats.append(feats)
        trainval_feats = pd.DataFrame(trainval_feats)
        test_feats = []
        for s in tqdm(test.sampleplotid.unique()):
            feats = process_image_features(f'{path}/{s}.tif', mask_plot, radius)
            feats['sampleplotid'] = s
            test_feats.append(feats)
        test_feats = pd.DataFrame(test_feats)
        trainval = trainval.merge(trainval_feats, on='sampleplotid', how='left')
        test = test.merge(test_feats, on='sampleplotid', how='left')
        
        if log_y:
            trainval[target_col] = np.log1p(trainval[target_col])
            test[target_col] = np.log1p(test[target_col])
        procs = None
        if normalize:
            procs = [Normalize]#.from_stats(*norm_stats)]
        feature_cols = [k for k in trainval_feats.columns if k != 'sampleplotid']
        trainval_tb = TabularPandas(trainval, procs=procs,
                                    cont_names=feature_cols, y_names=target_col,
                                    splits=ColSplitter(col='is_valid')(trainval))
        test_tb = TabularPandas(test, procs=procs,
                                cont_names=feature_cols, y_names=target_col)
        if save_path:
            trainval.to_csv(f'{save_path}/image_trainval.csv', index=False)
            test.to_csv(f'{save_path}/image_test.csv', index=False)
            with open(f'{save_path}/image_features.txt', 'w') as f:
                f.writelines("%s\n" % c for c in feature_cols)
        return trainval_tb, test_tb
    
    def load_image(self, path, target_col, normalize:bool=True, log_y:bool=False) -> Tuple[TabularPandas, TabularPandas]:
        "Load previously preprocessed image data"
        trainval = pd.read_csv(f'{path}/image_trainval.csv')
        test = pd.read_csv(f'{path}/image_test.csv')
        with open(f'{path}/image_features.txt', 'r') as f:
            feature_cols = [c.rstrip() for c in f.readlines()]
    
        if log_y:
            trainval[target_col] = np.log1p(trainval[target_col])
            test[target_col] = np.log1p(test[target_col])
        procs = None
        if normalize:
            procs = [Normalize]#.from_stats(*norm_stats)]
        trainval_tb = TabularPandas(trainval, procs=procs,
                                    cont_names=feature_cols, y_names=target_col,
                                    splits=ColSplitter(col='is_valid')(trainval))
        test_tb = TabularPandas(test, procs=procs,
                                cont_names=feature_cols, y_names=target_col)
        return trainval_tb, test_tb
    
    def preprocess(self, target_col, path, lidar_pref, image_pref, min_h:float=1.5,  
                   mask_plot:bool=True, normalize:bool=True, log_y:bool=False, 
                   save_path:str=None) -> Tuple[TabularPandas, TabularPandas]:
        "Preprocess dataframes and return (train_val, test) -tuple"
        trainval = self.train_val_df.copy()
        test = self.test_df.copy()
        feature_cols = point_cloud_metric_cols
        trainval[point_cloud_metric_cols] = trainval.progress_apply(lambda row: point_cloud_metrics(f'{path}/{lidar_pref}/{row.sampleplotid}.las',
                                                                                                    row.x, row.y, 
                                                                                                    min_h=min_h, 
                                                                                                    mask_plot=mask_plot), 
                                                                    axis=1, result_type='expand')
        test[point_cloud_metric_cols] = test.progress_apply(lambda row: point_cloud_metrics(f'{path}/{lidar_pref}/{row.sampleplotid}.las',
                                                                                            row.x, row.y, 
                                                                                            min_h=min_h, 
                                                                                            mask_plot=mask_plot), 
                                                            axis=1, result_type='expand')

        trainval_feats = []
        for s in tqdm(trainval.sampleplotid.unique()):
            feats = process_image_features(f'{path}/{image_pref}/{s}.tif', mask_plot, radius=31)
            feats['sampleplotid'] = s
            trainval_feats.append(feats)
        trainval_feats = pd.DataFrame(trainval_feats)
        test_feats = []
        for s in tqdm(test.sampleplotid.unique()):
            feats = process_image_features(f'{path}/{image_pref}/{s}.tif', mask_plot, radius=31)
            feats['sampleplotid'] = s
            test_feats.append(feats)
        test_feats = pd.DataFrame(test_feats)
        trainval = trainval.merge(trainval_feats, on='sampleplotid', how='left')
        test = test.merge(test_feats, on='sampleplotid', how='left')
        
        feature_cols = feature_cols + [c for c in trainval_feats.columns if c != 'sampleplotid']
        
        if log_y:
            trainval[target_col] = np.log1p(trainval[target_col])
            test[target_col] = np.log1p(test[target_col])
        procs = None
        if normalize:
            procs = [Normalize]#.from_stats(*norm_stats)]
        trainval_tb = TabularPandas(trainval, procs=procs,
                                    cont_names=feature_cols, y_names=target_col,
                                    splits=ColSplitter(col='is_valid')(trainval))
        test_tb = TabularPandas(test, procs=procs,
                                cont_names=feature_cols, y_names=target_col)
        
        if save_path:
            trainval.to_csv(f'{save_path}/las_image_trainval.csv', index=False)
            test.to_csv(f'{save_path}/las_image_test.csv', index=False)
            with open(f'{save_path}/las_image_features.txt', 'w') as f:
                f.writelines("%s\n" % c for c in feature_cols)
        return trainval_tb, test_tb
    
    def load_las_image(self, path, target_col, normalize:bool=True, log_y:bool=False) -> Tuple[TabularPandas, TabularPandas]:
        "Load previously preprocessed image data"
        trainval = pd.read_csv(f'{path}/las_image_trainval.csv')
        test = pd.read_csv(f'{path}/las_image_test.csv')
        with open(f'{path}/las_image_features.txt', 'r') as f:
            feature_cols = [c.rstrip() for c in f.readlines()]
    
        if log_y:
            trainval[target_col] = np.log1p(trainval[target_col])
            test[target_col] = np.log1p(test[target_col])
        procs = None
        if normalize:
            procs = [Normalize]#.from_stats(*norm_stats)]
        trainval_tb = TabularPandas(trainval, procs=procs,
                                    cont_names=feature_cols, y_names=target_col,
                                    splits=ColSplitter(col='is_valid')(trainval))
        test_tb = TabularPandas(test, procs=procs,
                                cont_names=feature_cols, y_names=target_col)
        return trainval_tb, test_tb 
    


# +
# export

def process_one(path, plot_x:float=None, plot_y:float=None, mask_plot=True, radius=9) -> pd.Series:
    "Utility for predicting single point cloud"
    metrics = point_cloud_metrics(path, plot_x, plot_y, mask_plot, radius)
    df = pd.DataFrame(columns=point_cloud_metric_cols, data=[metrics])
    return df.iloc[0]


# -

# hide
def read_image_features_from_csv(self, fn, target_col, 
                                 normalize:bool=True, log_y:bool=False) -> Tuple[TabularPandas, TabularPandas]:
    trainval = self.train_val_df.copy()
    test = self.test_df.copy()
    feats = pd.read_csv(fn)
    feature_cols = [c for c in feats.columns if c != 'sampleplotid']
    trainval_feats = feats[feats.sampleplotid.isin(trainval.sampleplotid.unique())]
    test_feats = feats[feats.sampleplotid.isin(test.sampleplotid.unique())]
    trainval = trainval.merge(trainval_feats, on='sampleplotid', how='left')
    test = test.merge(test_feats, on='sampleplotid', how='left')
    if log_y:
        trainval[target_col] = np.log1p(trainval[target_col])
        test[target_col] = np.log1p(test[target_col])
    procs = None
    if normalize:
        procs = [Normalize]#.from_stats(*norm_stats)]
    trainval_tb = TabularPandas(trainval, procs=procs,
                                cont_names=feature_cols, y_names=target_col,
                                splits=ColSplitter(col='is_valid')(trainval))
    test_tb = TabularPandas(test, procs=procs,
                            cont_names=feature_cols, y_names=target_col)

    return trainval_tb, test_tb


# +
# hide

from nbdev.export import notebook2script
notebook2script()
# !nbdev_build_docs
# -



StarcoderdataJupyterScriptsDedupFiltered