#!/usr/bin/python3
"""
    Copyright (C) 2023 NIBIO <https://www.nibio.no/>. 

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as
    published by the Free Software Foundation, either version 3 of the
    License, or (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""

import subprocess, os
from dotenv import load_dotenv
from datetime import datetime
import pytz
import netCDF4 as nc
from jinja2 import Environment, FileSystemLoader

load_dotenv()

weatherdata_path="in/"
tmp_path="tmp/"
out_path="out/"

# TODO: Make this truly independent of time zones/Norwegian conditions?
utc_offset = "+02:00"
local_timezone = pytz.timezone("Europe/Oslo")


# Calculate cumulated degree days above 5 degrees after 1st of April
# Remove all values before April 1st
subprocess.run(f"cdo -selname,TM -seldate,2023-04-01T00:00:00,2023-12-31T00:00:00 {weatherdata_path}met_1_0km_nordic-2023.nc {tmp_path}TM_from_april.nc", shell=True)
# Subtracting 5 deg C from all cells
subprocess.run(f"cdo -subc,5 {tmp_path}TM_from_april.nc {tmp_path}TM_minus_5.nc", shell=True)
# Create an .nc file with all cells containing 1 if cell in tgminus5.nc is >=0, 0 otherwise
subprocess.run(f"cdo -gtc,0 {tmp_path}TM_minus_5.nc {tmp_path}TM_minus_5_gtc.nc", shell=True)

# Multiplying tgminus5.nc with tgminus5gtc.nc, so that all sub zero values are set to 0
subprocess.run(f"cdo -mul {tmp_path}TM_minus_5.nc {tmp_path}TM_minus_5_gtc.nc {tmp_path}TM_minus_5_nozero.nc", shell=True)
# Accumulate day degrees with the corrected values
subprocess.run(f"cdo -timcumsum {tmp_path}TM_minus_5_nozero.nc {tmp_path}DD_unmasked.nc", shell=True)



# Mask the output with Norway land borders
subprocess.run(f"cdo -maskregion,Norge_landomrader.csv {tmp_path}DD_unmasked.nc {tmp_path}DD.nc", shell=True)

# Add the DD threshold classification => warning status
# (A>0)*(A<=260)*2 + (A>260)*(A<=360)*3 + (A>360)*(A<=560)*4 + (A>560)*0
subprocess.run(f'cdo -aexpr,"WARNING_STATUS=(TM>0)*(TM<=260)*2 + (TM>260)*(TM<=360)*3 + (TM>360)*(TM<=560)*4 + (TM>560)*0" {tmp_path}DD.nc {tmp_path}result.nc', shell=True)

# Split the combined file into daily .nc files again, with YYYY-MM-DD in the filename. Convert to corresponding GeoTIFF files
# Variables that needs discrete classification, must be integers in order for mapserver to work properly (Don't ask!)
# Since we need WARNING_STATUS to be discretely classified, we need to create a separate GeoTIFF file for it
result = nc.Dataset(f'{tmp_path}result.nc', 'r')
timesteps = result.variables["time"][:]
#print(timesteps)
timestep_index = 1
timestep_dates = []
for timestep in timesteps:
    timestep_date = datetime.fromtimestamp(timestep).astimezone(local_timezone)
    file_date = timestep_date.astimezone(local_timezone).strftime("%Y-%m-%d")
    print(f"{timestep_index}: {file_date}")
    timestep_dates.append(file_date)
    # We must delete the GeoTIFF file before merging
    subprocess.run(f'rm {tmp_path}result_*{file_date}_lcc.tif', shell=True)
    subprocess.run(f'rm {out_path}result_*{file_date}.tif', shell=True)
    subprocess.run(f'gdal_translate -ot Int16 -of GTiff  -b {timestep_index} NETCDF:"{tmp_path}result.nc":WARNING_STATUS {tmp_path}result_WARNING_STATUS_{file_date}_lcc.tif', shell=True)
    subprocess.run(f'gdal_translate -ot Float32 -of GTiff -b {timestep_index} NETCDF:"{tmp_path}result.nc":TM {tmp_path}result_{file_date}_lcc.tif', shell=True)
    # Need to reproject the files, to ensure we have the projection given in the generted mapfile. We always use EPSG:4326 for this
    subprocess.run(f'gdalwarp -t_srs EPSG:4326 {tmp_path}result_WARNING_STATUS_{file_date}_lcc.tif {out_path}result_WARNING_STATUS_{file_date}.tif', shell=True)
    subprocess.run(f'gdalwarp -t_srs EPSG:4326 {tmp_path}result_{file_date}_lcc.tif {out_path}result_{file_date}.tif', shell=True)
    timestep_index = timestep_index + 1

if len(timestep_dates) != len(set(timestep_dates)):
    print("ERROR: Something is wrong with your timesteps")
    #print(timestep_dates)

print(len(timestep_dates))

# Generate mapfile
# The paths should be set in a .env file
env = Environment(loader=FileSystemLoader('.'))
template = env.get_template("mapfile/template.j2")
output = template.render({
     "timestep_dates": timestep_dates,
     "mapserver_data_dir": os.getenv("MAPSERVER_DATA_DIR"),
     "mapserver_mapfile_dir": os.getenv("MAPSERVER_MAPFILE_DIR"),
     "mapserver_log_file": os.getenv("MAPSERVER_LOG_FILE"),
     "mapserver_image_path": os.getenv("MAPSERVER_IMAGE_PATH"),
     "mapserver_extent": os.getenv("MAPSERVER_EXTENT")
})
mapfile_outdir = os.getenv("MAPFILE_DIR")
with open(f"{mapfile_outdir}/PSILARTEMP.map", 'w') as f:
        f.write(output)