# ApexRMS - October 2022
# External Python script to run external program between every SyncroSim timestep.
# Setup ------------------------------------------------------------------------

## Load libraries
import pysyncrosim as ps
import pandas as pd
import numpy as np
import rasterio
import os 

## Set working directory from SyncroSim environment
e = ps.environment._environment()
wrk_dir = e.temp_directory.item()
if os.path.exists(wrk_dir) == False:
    os.mkdir(wrk_dir)

## Set up environment variables
timestep = e.before_timestep.item()                       
prev_timestep = timestep - 1
input_datasheet_name = "stsim_OutputSpatialAge"
output_datasheet_name = "stsim_TransitionSpatialMultiplier"
iteration = e.before_iteration.item()
transition_group_name = "Fire [Type]"

## Set path to store modified TIF file once analysis is complete
temp_name = "inputRaster.it" + str(iteration) + ".ts" + str(timestep) + ".tif"
temp_file_path = os.path.join(wrk_dir, temp_name)

## Load Library, Result Scenario, and Parent Scenario
myLibrary = ps.library(e.library_filepath.item())
myScenario = myLibrary.scenarios(sid = e.scenario_id.item())
my_scenario_parentid = int(myScenario.parent_id)
my_parent_scenario = myLibrary.scenarios(sid = my_scenario_parentid)

## Load raster output from previous timestep and Result Scenario
raster_map = myScenario.datasheet_rasters(datasheet = input_datasheet_name,
                                          timestep = prev_timestep,
                                          iteration = iteration)
raster_values = raster_map.values()

# Spatial processing below, replace with your own analysis ---------------------

## Create function to find the indices of the highest value cells in raster
def find_oldest_cells(raster_values, num_cells):
    flat = raster_values.flatten()
    indices = np.argpartition(flat, -num_cells)[-num_cells:]
    indices = indices[np.argsort(-flat[indices])]
    return np.unravel_index(indices, raster_values.shape)

## Reclassify 20 oldest cells as 1 and all other cells as 0
oldest_cells = find_oldest_cells(raster_values, 20)

mask = np.ones_like(raster_values, bool)
mask[oldest_cells] = False

raster_values[oldest_cells] = 1
raster_values[mask] = 0

# Spatial processing completed, continue below with template -------------------

# Save results back to SyncroSim -----------------------------------------------
## Save the updated raster to previously defined filepath

with rasterio.open(
    temp_file_path,
    'w',
    driver='GTiff',
    height=raster_values.shape[0],
    width=raster_values.shape[1],
    count=1,
    dtype=raster_values.dtype,
    crs='+proj=latlong',
    ) as dst:
    dst.write(raster_values, 1)

## Update Datasheet with results from spatial processing
my_output_dataframe = pd.DataFrame({
    "Iteration": [iteration],
    "Timestep": [timestep],
    "TransitionGroupID": [transition_group_name], 
    "MultiplierFileName": [temp_file_path]})

## Save the modified Datasheet to the Parent Scenario
my_parent_scenario.save_datasheet(name = output_datasheet_name, data = my_output_dataframe, append = True)