#!/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
import os
import sys
from dotenv import load_dotenv
from datetime import datetime
import pytz
import netCDF4 as nc
import configparser
from jinja2 import Environment, FileSystemLoader

# Load deployment specific settings
load_dotenv()
# Get language stuff
config = configparser.ConfigParser()
config.read("PSILARTEMP.cfg")

today = datetime.now().date()
if len(sys.argv) > 1:
    year = int(sys.argv[1])
else:
    year = today.year

recurring_start_date = os.getenv("RECURRING_START_DATE")
weatherdata_path = f'{os.getenv("WEATHER_DATA_DIR")}{year}/'
tmp_path = "tmp/"

out_path = f'{os.getenv("DATA_DIR")}{year}/'
os.makedirs(out_path, exist_ok=True)

# TODO: Put timezone out in .env
local_timezone = pytz.timezone("Europe/Oslo")

# Ensure that model is not run for future year
if year > today.year:
    print(f"Cannot run model for future year ({year}). Quit.")
    exit(0)

# If we are before the recurring start date, exit nicely
start_time = datetime.strptime(f"{year}-{recurring_start_date}","%Y-%m-%d")
if datetime.now() <= start_time:
    print(f"Today is before the configured start time of {start_time}. Exiting.")
    exit(0)
    
# """
# Calculate cumulated degree days above 5 degrees after 1st of March
# Remove all values before March 1st
subprocess.run(
    f"cdo -selname,TM -seldate,{year}-{recurring_start_date}T00:00:00,{year}-12-31T00:00:00 {weatherdata_path}met_1_0km_nordic-{year}.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 results using a CSV file with polygons
# Env variable MASK_FILE must be set
if os.getenv("MASK_FILE") is not None:
    mask_file = os.getenv("MASK_FILE")
    print(f"Applying mask file {mask_file} to result.nc")
    subprocess.run(
        f"cdo -maskregion,{mask_file} {tmp_path}DD_unmasked.nc {tmp_path}DD.nc",
        shell=True,
    )
else:
    os.rename(f"{tmp_path}DD_unmasked.nc", f"{tmp_path}DD.nc")

# Add the DD threshold classification => warning status
# 0: DD == 0            # Forecast not started
# 2: DD <= 260          # Flight period not started
# 3: 260 < DD <= 360    # Flight period starting
# 4: 360 < DD <= 560    # Peak flight period
# 5: 560 < DD           # 1st generation flight period ended
# (A>0)*(A<=260)*2 + (A>260)*(A<=360)*3 + (A>360)*(A<=560)*4 + (A>560)*5
subprocess.run(
    f'cdo -aexpr,"WARNING_STATUS=(TM>0)*(TM<=260)*2 + (TM>260)*(TM<=360)*3 + (TM>360)*(TM<=560)*4 + (TM>560)*5" {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)
    with open("/dev/null", "w") as devnull:
        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,
            stdout=devnull,
        )
        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,
            stdout=devnull,
        )
        # 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,
            stdout=devnull,
        )
        subprocess.run(
            f"gdalwarp -t_srs EPSG:4326 {tmp_path}result_{file_date}_lcc.tif {out_path}result_{file_date}.tif",
            shell=True,
            stdout=devnull,
        )
    # """
    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
# Building data sets for language specific legends
languages = []
language_codes = config["i18n"]["languages"].split(",")
for language_code in language_codes:
    language = {"language_code": language_code}
    if ("i18n.%s" % language_code) in config:
        for keyword in config["i18n.%s" % language_code]:
            language[keyword] = config["i18n.%s" % language_code][keyword]
        languages.append(language)

# 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"),
        "languages": languages,
        "language_codes": language_codes,
    }
)
mapfile_outdir = f'{os.getenv("MAPFILE_DIR")}{year}/'
os.makedirs(mapfile_outdir, exist_ok=True)
with open(f"{mapfile_outdir}PSILARTEMP.map", "w") as f:
    f.write(output)