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

# + [markdown] tags=[]
# # 24. Lidar point cloud cluster processing
# -

# This notebooks sets up files to process all Lidar point clouds on the Jassmin cluster

# + tags=[]
# Convenient jupyter setup
# %load_ext autoreload
# %autoreload 2

# + tags=[]
import geopandas as gpd
import pandas as pd
import pathlib
import datetime
from tqdm.autonotebook import tqdm

from src.constants import USER_PATH, PAISAGENSLIDAR_PATH, EBALIDAR_PATH, SRC_PATH
from src.utils.os import list_content, latest_matching_file
from src.utils.slurm import sbatch_header, sbatch

from src.processing.slurm_jobcreator import lidR_gridmetrics_job, lidR_processing_job

# + [markdown] tags=[]
# ### Helper functions

# + tags=[]
is_empty = lambda folder: not any(folder.iterdir())
contains = lambda folder, name: len(list(folder.glob(name))) > 0

# + [markdown] tags=[]
# ## 1. Paisagenslidar - Cluster normalisation

# + tags=[]
paisagens_meta_path = latest_matching_file(PAISAGENSLIDAR_PATH / "*metadata")
paisagens_meta = gpd.read_file(paisagens_meta_path)

# + tags=[]
paisagens_paths = paisagens_meta.groupby("survey").first().path.apply(lambda x: pathlib.Path(x).parent)
paisagens_years = paisagens_meta.groupby("survey")[["year", "crs", "region", "plot_code", "creation_date"]].first()
paisagens_nfiles = paisagens_meta.groupby("survey").filename.count()
paisagens_size = paisagens_meta.groupby("survey")[["area (ha)",  "file_size (MB)", "points (Mio.)"]].sum()

# + tags=[]
paisagens_surveys = pd.concat([paisagens_paths, paisagens_years, paisagens_nfiles, paisagens_size], join="inner", axis=1)
paisagens_surveys["name"] = paisagens_surveys.apply(lambda x: f"{x.region}_{x.plot_code}_{x.year}", axis=1)
paisagens_surveys = paisagens_surveys.sort_values(by=["year", "region", "plot_code"])
# -

# ### Normalisation: Create slurm job

# + tags=[]
for idx, survey in tqdm(paisagens_surveys.iterrows(), total=len(paisagens_surveys)):
    
    out_dir = PAISAGENSLIDAR_PATH / "processed" / str(survey["year"]) / survey["name"]
    
    if contains(out_dir, "processed") and not is_empty(out_dir / "processed"):
        continue
    
    job_file = lidR_processing_job(survey, 
                                   out_dir = out_dir,
                                   queue="short-serial",
                                   max_time = datetime.time(23, 59, 0),
                                   chunk_size=0,
                                   buffer=50,
                                   filter_quantile = 0.95,
                                   filter_sensitivity = 1.1,
                                   filter_gridsize= 10,
                                   as_file=False, 
                                   perform_check=False,
                                   compress_intermediates=True, 
                                   save_intermediates=True,
                                   required_memory_per_node=15000,
                                  )
        
    print(job_file)
    break
    # Submit job
    #sbatch(job_file)
    print(f"Submitted job for {survey['name']}")
# -

# ### Metrics: Create slurm job

for idx, survey in paisagens_surveys.iterrows():
    
    if survey["path"] is None or pd.isna(survey["path"]):
        continue
    
    out_dir = PAISAGENSLIDAR_PATH / "processed" / str(survey["year"]) / survey["name"]
    
    if contains(out_dir, "processed") and not is_empty(out_dir / "processed"):
        pass
    else:
        print(survey["name"], "not available.")
        continue
    
    job_file = lidR_gridmetrics_job(survey, 
                                    out_dir,
                                    metrics = {"n_points": [1],
                                               "n_pulses": [1],
                                               "n_ground_points": [1],
                                               "max": [10, 50],
                                               "quantile_0.05": [10, 50],
                                               "quantile_0.1": [10, 50],
                                               "kurtosis": [10, 50],
                                               "mean": [10, 50],
                                               "interquartile_range": [10, 50],
                                               "longo_biomass": [10, 50], 
                                               "canopy_height": [1]

                                    },
                                    buffer = 10,
                                    chunk_size = 0,
                                    as_file = False,
                                    overwrite=False,
                                    queue="short-serial",
                                    required_memory_per_node=8000,    
                                    max_time = datetime.time(10, 0, 0)
                                   )
    print(job_file)
    break
    # Submit job
    #sbatch(job_file)
    print(f"Submitted job for {survey['name']}")

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

# + [markdown] tags=[] toc-hr-collapsed=true
# ## 2. EBA - Cluster normalisation

# + tags=[]
eba_meta_path = latest_matching_file(EBALIDAR_PATH / "*metadata")
eba_meta = gpd.read_file(eba_meta_path)

# + tags=[]
eba_laz = {path.name: {"path": path,
                       "name": path.name.split(".")[0],
                       "size": round(path.stat().st_size / 1024**2,2)
                      } 
           for path in (EBALIDAR_PATH / "laz_EBA").glob("*.laz")
          }

# + tags=[]
df = pd.DataFrame(eba_laz).T.convert_dtypes()
df = df.rename(columns={"size": "size (MB)"})
df["reduced_filename"] = df["name"] + ".laz"

# + tags=[]
eba_new_meta = pd.merge(eba_meta.drop(columns=["path", "size (MB)", "name"]), df, how="outer", on="reduced_filename")
eba_new_meta["size (MB)"] = eba_new_meta["size (MB)"].astype(float)
eba_new_meta.index = eba_new_meta["reduced_filename"]
eba_new_meta.index.name = "survey"
eba_new_meta = eba_new_meta.sort_index()
eba_new_meta

# + tags=[]
#eba_new_meta.to_file(eba_meta_path.parent / f"eba_file_metadata_v{datetime.date.today()}", driver="GPKG")

# + tags=[]
smallest = df.iloc[df["size (MB)"].values.argmin():df["size (MB)"].values.argmin()+1]
smallest
# -

# ### Normalisation: Create slurm job

# + tags=[]
for idx, survey in tqdm(eba_new_meta.iterrows(), total=len(eba_new_meta)):
    
    if survey["path"] is None or pd.isna(survey["path"]):
        continue
    
    out_dir = EBALIDAR_PATH / "laz_EBA_processed" / survey["name"]
    
    if contains(out_dir, "processed") and not is_empty(out_dir / "processed"):
        continue
    #else:
        #print("Not processed:", survey["name"])
        #continue
    
    job_file = lidR_processing_job(survey, 
                                   out_dir = out_dir,
                                   queue="short-serial",
                                   max_time = datetime.time(23, 59, 0),
                                   chunk_size=1000,
                                   buffer=50,
                                   filter_quantile = 0.95,
                                   filter_sensitivity = 1.1,
                                   filter_gridsize= 10,
                                   as_file=False, 
                                   perform_check=False,
                                   compress_intermediates=True, 
                                   save_intermediates=True,
                                   required_memory_per_node=16000,
                                  )
    
    print(job_file)
    break
    # Submit job
    sbatch(job_file)
    print(f"Submitted job for {survey['name']}")
# -

# ### Metrics: Create slurm job

# + tags=[]
for idx, survey in tqdm(eba_new_meta.iterrows(), total=len(eba_new_meta)):
    
    if survey["path"] is None or pd.isna(survey["path"]):
        continue
    
    out_dir = EBALIDAR_PATH / "laz_EBA_processed" / survey["name"]
    
    job_file = lidR_gridmetrics_job(# File params
                                    survey, 
                                    out_dir,

                                    # Processing params
                                    metrics = {"point_density": [1],
                                               "pulse_density": [1],
                                               "ground_point_density": [1],
                                               "n_points": [1],
                                               "n_pulses": [1],
                                               "n_ground_points": [1],
                                               "max": [10, 50],
                                               "quantile_0.05": [10, 50],
                                               "quantile_0.1": [10, 50],
                                               "kurtosis": [10, 50],
                                               "mean": [10, 50],
                                               "interquartile_range": [10, 50],
                                               "longo_biomass": [10, 50], 
                                               "canopy_height": [1, 2]

                                    },
                                    buffer = 10,
                                    chunk_size = 0,
                                    as_file = False,

                                    # Queue params
                                    queue="short-serial",
                                    max_time = datetime.time(10, 0, 0),
                                    #expected_time = datetime.time(1, 30, 0),
                                    required_memory_per_node=8000,
                               )
    print(job_file)
    # Submit job
    break
    sbatch(job_file)
    print(f"Submitted job for {survey['name']}")
# -

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

# + [markdown] tags=[]
# # To be deleted: Expedited processing for relevant files

# + tags=[]
relevant = [
    "NP_T-0408.laz",
    "NP_T-0859.laz",
    "NP_T-0145.laz",
    "NP_T-0751.laz",
    "NP_T-2002.laz",
    "NP_T-0416.laz",
    "NP_T-0669.laz",
]

for idx, survey in eba_new_meta.loc[relevant].iterrows():
    
    #if survey["path"] is None or pd.isna(survey["path"]):
    #    continue
    
    out_dir = EBALIDAR_PATH / "laz_EBA_processed" / survey["name"]
    
    job_file = create_lidar_gridmetrics_job(idx, 
                                survey, 
                                out_dir,
                                metrics = {"point_density": [1],
                                           "pulse_density": [1],
                                           "ground_point_density": [1],
                                           "n_points": [1],
                                           "n_pulses": [1],
                                           "n_ground_points": [1],
                                           "max": [10, 50],
                                           "quantile_0.05": [10, 50],
                                           "quantile_0.1": [10, 50],
                                           "kurtosis": [10, 50],
                                           "mean": [10, 50],
                                           "interquartile_range": [10, 50],
                                           "longo_biomass": [10, 50], 
                                           "canopy_height": [1, 2]
                                
                                },
                                buffer = 10,
                                chunk_size = 0,
                                as_file = True,
                                overwrite=False,
                                queue="test",
                                
                               )
    print(job_file)
    break
    # Submit job
    #sbatch(job_file)
    print(f"Submitted job for {survey['name']}")

# + tags=[]
paisagens_relevant = [
    "JAM_A03_2013_laz",
    "JAM_A02_2013_laz",
    "JAM_A02a_2014_laz",
    "JAM_A02c_2015_LiDAR",
    "SAN_A02_2014_laz",
    
]

for idx, survey in paisagens_surveys.loc[paisagens_relevant][4::].iterrows():
    
    #if survey["path"] is None or pd.isna(survey["path"]):
    #    continue
    
    out_dir = PAISAGENSLIDAR_PATH / "processed" / str(survey["year"]) / survey["name"]
    
    if contains(out_dir, "processed") and not is_empty(out_dir / "processed"):
        pass
    else:
        print(survey["name"], "not available.")
        continue
    
    job_file = lidR_gridmetrics_job(survey, 
                                    out_dir,
                                    metrics = {"n_points": [1],
                                               "n_pulses": [1],
                                               "n_ground_points": [1],
                                               "max": [10, 50],
                                               "quantile_0.05": [10, 50],
                                               "quantile_0.1": [10, 50],
                                               "kurtosis": [10, 50],
                                               "mean": [10, 50],
                                               "interquartile_range": [10, 50],
                                               "longo_biomass": [10, 50], 
                                               "canopy_height": [1]

                                    },
                                    buffer = 10,
                                    chunk_size = 0,
                                    as_file = False,
                                    overwrite=False,
                                    queue="test",
                                    required_memory_per_node=8000,    
                                    max_time = datetime.time(10, 0, 0)
                                   )
    print(job_file)
    break
    # Submit job
    sbatch(job_file)
    print(f"Submitted job for {survey['name']}")


# -

# ----

# # Legacy

# + jupyter={"source_hidden": true} tags=[]
def create_lidar_processing_job(idx, 
                                survey, 
                                r_file,
                                out_dir,
                                buffer = 50,
                                chunk_size = 0,
                                filter_quantile = 0.95, 
                                filter_sensitivity = 1.1,
                                filter_gridsize = 10,
                                compress_intermediates = True,
                                save_intermediates = True,
                                perform_check = True,
                                as_file = True,
                                **kwargs
                               ):
    # Path for lidar data
    lidar_path = pathlib.Path(survey["path"])
    assert lidar_path.exists()
    
    # Set up save prefix
    save_prefix = survey["name"] + "_"

    # Create working directory path
    slurm_dir = out_dir / "slurm"
    slurm_dir.mkdir(parents=True, exist_ok=True)

    # Create sbatch header
    header = sbatch_header(
                  job_name = idx + "-->" + survey["name"],
                  working_directory=slurm_dir,
                  output_file_pattern=f"slurm_{survey['name']}-%j.out",
                  error_file_pattern=f"slurm_{survey['name']}-%j.err",
                  **kwargs)

    body = seperator("Activating conda environment")
    body += "source ~/.bashrc\n"
    body += "conda activate r4_env\n"
    body += seperator("Checking R version")
    body += echo("$(which R)")
    #body += "echo $(which R)\n"
    body += echo("$(R --version)")
    body += seperator("Analysing file")
    body += echo(f"template generation time: {datetime.datetime.now()}")
    body += echo(f"survey: {survey['name']}")
    body += echo(f"path: {lidar_path}")
    body += seperator("Executing R script")
    body += (f"Rscript --verbose --no-save {r_file}" 
             f" --lidar_path={lidar_path}" 
             f" --save_path={out_dir}"
             f" --buffer={buffer}"
             f" --chunk_size={chunk_size}"
             f" --filter_quantile={filter_quantile}"
             f" --filter_sensitivity={filter_sensitivity}"
             f" --filter_gridsize={filter_gridsize}"
             f" --save_prefix={save_prefix}"
             f" --perform_check={'TRUE' if perform_check else 'FALSE'}"
             f" --save_intermediates={'TRUE' if save_intermediates else 'FALSE'}"
             f" --compress_intermediates={'TRUE' if compress_intermediates else 'FALSE'}"
            )

    if as_file:
        job_file = slurm_dir / f"slurm_{survey['name']}.sh"
        with open(job_file, "w") as file:
            file.write(header + body)
        return job_file
    else: 
        return header+body


# + jupyter={"source_hidden": true} tags=[]
ALLOWED_METRICS = ("point_density", "pulse_density", "ground_point_density",
                   "n_points", "n_pulses", "n_ground_points",
                   "max", "standard_dev", "mask", "mean", "kurtosis", 
                   "interquartile_range", "quantile", "longo_biomass")

r_gridmetrics_file = "/home/users/svm/Code/gedi_biomass_mapping/src/processing/lidar_gridmetrics_script.R"
r_canopy_height_file = "/home/users/svm/Code/gedi_biomass_mapping/src/processing/lidar_canopy_height_script.R"

def create_lidar_gridmetrics_job(idx, 
                                survey, 
                                out_dir,
                                metrics,
                                buffer = 10,
                                chunk_size = 0,
                                save_prefix="",
                                check_index=True,
                                overwrite=False,
                                as_file = False,
                                **kwargs
                               ):
    # Path for lidar data
    lidar_path = pathlib.Path(out_dir / "processed")
    assert lidar_path.exists()
    
    # Set up save prefix
    save_prefix = survey["name"] + "_"

    # Create working directory path
    slurm_dir = out_dir / "slurm"
    slurm_dir.mkdir(parents=True, exist_ok=True)
    grid_metrics_dir = out_dir / "grid_metrics"
    grid_metrics_dir.mkdir(parents=True, exist_ok=True)

    # Create sbatch header
    header = sbatch_header(
                  job_name = survey["name"] + "_grid_metrics",
                  working_directory=slurm_dir,
                  output_file_pattern=f"slurm_{survey['name']}_metrics-%j.out",
                  error_file_pattern=f"slurm_{survey['name']}_metrics-%j.err",
                  **kwargs)

    body = seperator("Activating conda environment")
    body += "source ~/.bashrc\n"
    body += "conda activate r4_env\n"
    body += seperator("Checking R version")
    body += echo("$(which R)")
    body += echo("$(R --version)")
    body += seperator("Analysing file")
    body += echo(f"template generation time: {datetime.datetime.now()}")
    body += echo(f"survey: {survey['name']}")
    body += echo(f"path: {lidar_path}")
    body += echo(f"metrics: {metrics}")
    body += seperator("Executing R script")
    
    if "canopy_height" in metrics.keys():
        for gridsize in metrics.pop("canopy_height"):
            body += f"\n# `Canopy height` at grid size `{gridsize}m`\n"
            body += (f"Rscript --verbose --no-save {r_canopy_height_file}" 
                     f" --lidar_path={lidar_path}"
                     f" --gridsize={gridsize}" 
                     f" --save_path={out_dir}"
                     f" --buffer={buffer}"
                     f" --chunk_size={chunk_size}"
                     f" --save_prefix={save_prefix}"
                     f" --overwrite={'TRUE' if overwrite else 'FALSE'}"
                     f" --check_index={'TRUE' if check_index else 'FALSE'}"
                     "\n"
                    )
    
    for metric in metrics.keys():
        assert metric in ALLOWED_METRICS or "quantile_" in metric, f"{metric} not allowed."
        for gridsize in metrics[metric]:
            body += f"\n# Metric `{metric}` at grid size `{gridsize}m`\n"
            body += (f"Rscript --verbose --no-save {r_gridmetrics_file}" 
                     f" --lidar_path={lidar_path}" 
                     f" --metric={metric if 'quantile' not in metric else 'quantile'}"
                     f" --gridsize={gridsize}" 
                     f" --save_path={out_dir}"
                     f" --buffer={buffer}"
                     f" --chunk_size={chunk_size}"
                     f" --save_prefix={save_prefix}"
                     f" --overwrite={'TRUE' if overwrite else 'FALSE'}"
                     f" --check_index={'TRUE' if check_index else 'FALSE'}"
                     f"{' --quantile='+str(metric.split('_')[-1]) if 'quantile' in metric else ''}"
                     "\n"
                    )

    if as_file:
        job_file = slurm_dir / f"slurm_{survey['name']}_grid_metrics.sh"
        with open(job_file, "w") as file:
            file.write(header + body)
        return job_file
    else: 
        return header+body

StarcoderdataJupyterScriptsDedupFiltered