#!/usr/bin/env python3
#  -*- coding: utf-8 -*-

"""
Runs 5 days backtracking simulations of oil particles from different sites in Qatar

Author: Thomas Anselain, Earth and Life Institute, UCLouvain, Belgium
Last modified: 7 April 2022
"""
#%% Import packages

import sys
from datetime import timedelta
from datetime import datetime
from opendrift.readers import reader_netCDF_CF_generic
from opendrift.models.openoil import OpenOil
import os
from netCDF4 import Dataset
from mesh_to_landmask import h5_to_landmask
import calendar
import numpy as np
import geopandas as gpd

basedir = "/export/miro/students/tanselain/OpenOil/"

abbr_to_num = {name: num for num, name in enumerate(calendar.month_abbr) if num}

#%% Beginning of loops
def OpenOil_run(source, mois, year, out) :
    path = basedir + f'Simulation_backward/Shape_desalination_plants/Point_station_buffer_decoupe_point_buffer_{source}.shp'
    df1 = gpd.read_file(path)
    x = len(df1)
    print(x)
    
    num =abbr_to_num[mois] 
    # Take into account the monthly averaged temperature
    b = {'Jan': 22.1, 'Feb': 20.85, 'Mar': 22.25, 'Apr': 24.85, 'May': 25.55, 'Jun': 31.25, 'Jul': 32.6, 'Aug': 33.3, 'Sep': 32.55, 'Oct': 30.9, 'Nov': 28.05, 'Dec': 24.35}
    temp = b[mois]
    
    # Number of loop in the month
    tstart = calendar.timegm([year_int, num, 1, 0, 0, 0])
    tend = calendar.timegm([year_int+1 if num == 12 else year_int, num%12+1, 1, 0, 0, 0])
    ndays = int((tend-tstart)/86400)
    print(mois, ndays)
    
    if mois == 'Sep' :
        nb =16
    else :
        nb = 1
        
    # Simulation loop
    for date in range(nb, ndays+1) :
        #if source == 'umm_al_houl' and mois == 'Jan' : 
        #    continue 
        #if source == 'umm_al_houl' and mois == 'Feb' : 
        #    continue 
        #if source == 'umm_al_houl' and mois == 'Mar' : 
        #    continue 
        #if source == 'umm_al_houl' and mois == 'Apr' : 
        #    continue
        # Find forcing data
        forcdir = basedir + f'Simulation_backward'
        if year == '2016' :
          wind_file = forcdir + f"/Forcing_data/wind/wind_2015_2016.nc" 
        if year == '2017' or year == '2018' : 
          wind_file = forcdir + f"/Forcing_data/wind/wind_2016_2018.nc"
        if year == '2019' or year == '2020' :
          wind_file = forcdir + f"/Forcing_data/wind/wind_2018_2020.nc"
        wave_file = forcdir + f"/Forcing_data/waves/{year}/waves_{year}{num:02d}.nc"
        mercator_file    = forcdir + f"/Forcing_data/currents+tides/{year}/currents+tides_{year}{num:02d}_extended_2px.nc"
        bathy_file = forcdir + "/Forcing_data/bathymetry_persian_gulf.nc"
        
        def add_std_names(filename):
            # to avoid os error
            os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE"
            with Dataset(filename, "r+") as f:
                for var in f.variables.values():
                    if not hasattr(var, "standard_name"):
                        # ECMWF wind files
                        if var.name == "u10":
                            setattr(var, "standard_name", "x_wind")
                        elif var.name == "v10":
                            setattr(var, "standard_name", "y_wind")
        
        #def invert_ecmwf_wind(filename):
        #     with Dataset(filename, "r+") as f:
        #         latvar = f.variables["latitude"]
        #         lat = np.asarray(latvar[:])
        #         if lat[1] < lat[0]:
        #             latvar[:] = lat[::-1]
        #             uvar, vvar = f.variables["u10"], f.variables["v10"]
        #             for i in range(uvar.shape[0]):
        #                 uvar[i,:] = np.flipud(np.array(uvar[i]))
        #                 vvar[i,:] = np.flipud(np.array(vvar[i]))
        #                
        #invert_ecmwf_wind(wind_file)
        
        
        # add_std_names(wind_file)
        reader_wind = reader_netCDF_CF_generic.Reader(
            filename=wind_file,
            name="wind",
            proj4='+proj=latlong'
        )
          
        
        reader_waves = reader_netCDF_CF_generic.Reader(
            filename=wave_file,
            name='waves',
            proj4="+proj=latlong"
        )
        
        reader_mercator = reader_netCDF_CF_generic.Reader(
            filename=mercator_file,
            name="mercator",
            proj4='+proj=latlong'
        )
    
        reader_bathy = reader_netCDF_CF_generic.Reader(
            filename= bathy_file,
            name="bathy",
            proj4='+proj=latlong'
        )


#%% Backward simulations of 5 days in each days of a given month
    
    
        #if year == '2018' and mois == 'Apr' and source == 'ras_laffan' and date <= 3 :
            #continue 
        #if year == '2017' and mois == 'Oct' and source == 'umm_al_houl' and date == 25 :
            #continue 
        o = OpenOil(loglevel=20, weathering_model='noaa')
        o.add_reader([reader_mercator, reader_wind, reader_waves])
        
        o.set_config('general:coastline_action', 'stranding')
        o.set_config('environment:constant:sea_water_temperature', temp)
        o.set_config('environment:constant:sea_water_salinity', 40)
        
        # Seeding some oil particles on 
        oiltype= 'QATAR MARINE' #'ARABIAN HEAVY'  'QATAR MARINE, OIL & GAS' 'QATAR MARINE, PHILLIPS' 'QATAR MARINE'
        start_time=datetime(year_int,num,date)  #2020 ou 2019
        time_set = timedelta(hours=24)
        if mois == 'Dec' and date == 31 :
            time_set = timedelta(hours=23)
        o.seed_from_shapefile( path, number=5*x, layername=None,  #de base 54 et pas 16
                                featurenum= None, time=start_time+time_set, oil_type=oiltype) #Time begin the next day at midnight
        
        # Adjusting some configuration
        o.set_config('processes:dispersion', True)
        o.set_config('processes:evaporation', True)
        o.set_config('processes:emulsification', True)
        o.set_config('drift:vertical_mixing', True)
        o.set_config('vertical_mixing:timestep', 2)

        # Running model
        o.run(steps=5*24*4, time_step=-900, time_step_output=900, outfile=out + f"out_{mois2}_{date}.nc")
                     
        # Print and plot results (additionnal)
        #o.plot_oil_budget()
        #o.plot(fast=False, filename= out + f"out_{mois2}_{date}.png", show=False)
        #o.animation(fast=True) 
        

#%% Give combinations station/month/year assessed

year = sys.argv[1] #(2016, 2017, 2018, 2019, 2020)
year_int = int(year)

#list_source = ['ras_laffan', 'abu_fontas', 'umm_al_houl', 'ras_laffan_port']
list_source = ['ras_laffan_port']

list_mois = ['Jan','Feb','Mar','Apr','May','Jun','Jul','Aug','Sep','Oct','Nov','Dec']
#list_mois = ['Sep','Oct','Nov','Dec']


for source in list_source :
    for mois in list_mois :
        a = {'Jan': 'january', 'Feb': 'february', 'Mar': 'march', 'Apr': 'april', 'May': 'may', 'Jun': 'june', 'Jul': 'july', 'Aug': 'august', 'Sep': 'september', 'Oct': 'october', 'Nov': 'november', 'Dec': 'december'}
        mois2 = a[mois]
        print(f'The month is {abbr_to_num[mois]:02d}')
        OpenOil_run(source, mois, year, basedir + f"Output_backward/{year}/{source}/{mois2}/")