#!/usr/bin/env Rscript

# FinnishForestAI — R workflow
# Usage:
#   Rscript r/forest_lidar_analysis.R data/raw/tile.laz outputs/r_run
# or use the included synthetic CSV.

args <- commandArgs(trailingOnly = TRUE)
input <- if (length(args) >= 1) args[[1]] else "data/demo/synthetic_forest_points.csv"
outdir <- if (length(args) >= 2) args[[2]] else "outputs/r_run"
dir.create(outdir, recursive = TRUE, showWarnings = FALSE)

required <- c("lidR", "terra", "ranger", "dplyr", "readr", "jsonlite", "ggplot2")
missing <- required[!vapply(required, requireNamespace, logical(1), quietly = TRUE)]
if (length(missing)) stop("Install packages first with: Rscript r/install_packages.R\nMissing: ", paste(missing, collapse=", "))

library(lidR)
library(terra)
library(ranger)
library(dplyr)
library(readr)
library(jsonlite)
library(ggplot2)

is_las <- grepl("\\.(las|laz)$", input, ignore.case = TRUE)

if (is_las) {
  las <- readLAS(input)
  if (is.empty(las)) stop("LAS/LAZ file contains no points")
  # NLS class 2 provides ground points. TIN normalization uses those ground returns.
  las_n <- normalize_height(las, tin())
  d <- as.data.frame(las_n@data)
  d$x <- las_n@data$X
  d$y <- las_n@data$Y
  d$z <- las_n@data$Z
  # After normalize_height(), Z is height above ground.
  d$height_agl <- pmax(d$z, 0)
  if (!"Intensity" %in% names(d)) d$Intensity <- 0
  if (!"ReturnNumber" %in% names(d)) d$ReturnNumber <- 1
  if (!"NumberOfReturns" %in% names(d)) d$NumberOfReturns <- 1
  if (!"Classification" %in% names(d)) d$Classification <- 1

  chm <- rasterize_canopy(las_n, res = 2, algorithm = p2r())
  writeRaster(chm, file.path(outdir, "canopy_height_model_R.tif"), overwrite = TRUE)
} else {
  d <- read_csv(input, show_col_types = FALSE)
  names(d) <- tolower(names(d))
  if (!all(c("x", "y", "z") %in% names(d))) stop("CSV needs x, y, z")
  if (!"classification" %in% names(d)) d$classification <- 1
  if (!"intensity" %in% names(d)) d$intensity <- 0
  if (!"return_number" %in% names(d)) d$return_number <- 1
  if (!"number_of_returns" %in% names(d)) d$number_of_returns <- 1

  ground <- d %>% filter(classification == 2)
  # Simple nearest-neighbour-like fallback for the portable CSV demo:
  baseline <- quantile(d$z, 0.02, na.rm = TRUE)
  d$height_agl <- pmax(d$z - baseline, 0)
  d$Intensity <- d$intensity
  d$ReturnNumber <- d$return_number
  d$NumberOfReturns <- d$number_of_returns
  d$Classification <- d$classification
}

veg <- d$height_agl[d$height_agl > 0.5]
canopy <- d$height_agl[d$height_agl > 2]
metrics <- list(
  points_used = nrow(d),
  canopy_cover_gt2m_fraction = mean(d$height_agl > 2, na.rm = TRUE),
  vegetation_height_mean_m = if(length(veg)) mean(veg, na.rm=TRUE) else NA,
  vegetation_height_p50_m = if(length(veg)) as.numeric(quantile(veg, .50, na.rm=TRUE)) else NA,
  vegetation_height_p95_m = if(length(veg)) as.numeric(quantile(veg, .95, na.rm=TRUE)) else NA,
  canopy_height_p95_m = if(length(canopy)) as.numeric(quantile(canopy, .95, na.rm=TRUE)) else NA,
  max_height_agl_m = max(d$height_agl, na.rm=TRUE)
)
write_json(metrics, file.path(outdir, "forest_metrics_R.json"), auto_unbox = TRUE, pretty = TRUE)
write_csv(as.data.frame(metrics), file.path(outdir, "forest_metrics_R.csv"))

# Baseline AI: learn useful NLS classes 2–5 where they exist.
train <- d %>%
  filter(Classification %in% c(2,3,4,5)) %>%
  transmute(
    height_agl = height_agl,
    intensity = Intensity,
    return_number = ReturnNumber,
    number_of_returns = NumberOfReturns,
    classification = factor(Classification)
  )

if (nrow(train) >= 500) {
  if (nrow(train) > 250000) train <- train[sample(seq_len(nrow(train)), 250000), ]
  set.seed(42)
  model <- ranger(
    classification ~ height_agl + intensity + return_number + number_of_returns,
    data = train,
    num.trees = 220,
    importance = "impurity",
    probability = TRUE
  )
  saveRDS(model, file.path(outdir, "vegetation_classifier_R.rds"))
  write_json(as.list(model$variable.importance), file.path(outdir, "variable_importance_R.json"), auto_unbox=TRUE, pretty=TRUE)
}

p <- ggplot(d[sample(seq_len(nrow(d)), min(30000, nrow(d))), ], aes(x = x, y = height_agl)) +
  geom_point(size = 0.2, alpha = 0.35) +
  labs(title = "LiDAR height profile", x = "Easting", y = "Height above ground (m)") +
  theme_minimal()
ggsave(file.path(outdir, "height_profile_R.png"), p, width=9, height=5, dpi=160)

cat("R analysis complete. Outputs:", outdir, "\n")