# This file is part of OpenDrift.
#
# OpenDrift is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, version 2
#
# OpenDrift 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with OpenDrift.  If not, see <https://www.gnu.org/licenses/>.
#
# Copyright 2015, 2020, Knut-Frode Dagestad, MET Norway

"""
This is a Pelagia noctiluca model based on the model_template and larvalfish 
model. 
This model uses currents for transport and introduces diel vertical migrations
as in Berline et al. 2013. It also uses environmental variables from MFS and biogeochemical 
variables from MedBFM to form an individual based model of Pelagia.
"""

import numpy as np
from opendrift.models.oceandrift import OceanDrift
from opendrift.models.oceandrift import Lagrangian3DArray
from suncalc import get_position
from datetime import datetime
import sys

class PelagiaElement(Lagrangian3DArray):
    """
    Extending Lagrangian3DArray with specific properties for jellyfsih
    """

    variables = Lagrangian3DArray.add_variables([
        ('biomass', {'dtype': np.float32,
                      'units': '',
                      'default': 10000.0}),
        ('p_production', {'dtype': np.float32,
                      'units': '',
                      'default': 0.0}),
        ('kb', {'dtype': np.float32,
                      'units': '',
                      'default': 0.0}),
        ('ks', {'dtype': np.float32,
                      'units': '',
                      'default': 0.0})])

class Pelagia_IBM(OceanDrift):
	
    ElementType = PelagiaElement
	
    required_variables = {
        'x_sea_water_velocity': {'fallback': 0},
        'y_sea_water_velocity': {'fallback': 0},
#        'sea_surface_wave_significant_height': {'fallback': 0},
 #       'x_wind': {'fallback': 0},
 #       'y_wind': {'fallback': 0},
        'land_binary_mask': {'fallback': None},
        'nearshore_mask': {'fallback': None},
        'sea_floor_depth_below_sea_level': {'fallback': 100},
#        'ocean_vertical_diffusivity': {'fallback': 0.01, 'profiles': True},
#        'ocean_mixed_layer_thickness': {'fallback': 50},
        'net_primary_production_of_biomass_expressed_as_carbon_per_unit_volume_in_sea_water': {'fallback': 5.0, 'profiles': True},
#        'sea_water_temperature': {'fallback': 10, 'profiles': True},
#        'sea_water_salinity': {'fallback': 34, 'profiles': True}
    }
    # use only values in defined depth range
    required_profiles_z_range = [0, -600]
   
    
    def __init__(self, *args, **kwargs):
        # Calling general constructor of parent class
        super(Pelagia_IBM, self).__init__(*args, **kwargs)
        # Configuration options
        self._add_config({
            'k1': {'type': 'float', 'default': 0,
                'min': 0, 'max': 10, 'units': 'day-1',
                'description': 'Growth parameter',
                'level': self.CONFIG_LEVEL_ESSENTIAL},
            'k2': {'type': 'float', 'default': 0,
                'min': 0, 'max': 100, 'units': 'mg m−3 day−1',
                'description': 'Half-saturation constant',
                'level': self.CONFIG_LEVEL_ESSENTIAL},
            'k3': {'type': 'float', 'default': 0,
                'min': 0, 'max': 10, 'units': 'day^-1',
                'description': 'Mortality',
                'level': self.CONFIG_LEVEL_ESSENTIAL},
           })
        
    
    def diel_vertical_migration(self):
        # Diel migration with upward / downward motion 
        # Max depth reached by Pelagia
        MaxDpth = -500.0
        # Vertical Swimming in m/s (2 m/min - Berline et al., 2013)
        VertSwimmingSpeed = 2/60
        # Minimum depth above bottom 
        MinAboveBottom = 5
        # Migration type
        MigrationType = 2  #(0 - no diel migrations, 1 - migration at fixed UTC, 2 - migration synchronized with Sun)
        # Particles close to the bottom should stay 5 m above the bottom, depth is negative
        Zmin = -1.*self.environment.sea_floor_depth_below_sea_level + MinAboveBottom
         
        if MigrationType == 0:
            print('No diel vertical migrations ...')
            
        # Old migration - using fixed time
        elif MigrationType == 1:
            print('UTC migrations ...')
            if self.time.hour > 7.0 and self.time.hour < 19.0:
                # Going down
                in_ocean = np.where(self.elements.z<0)[0]
                if len(in_ocean) > 0:
                    self.elements.z[in_ocean] = np.maximum(np.maximum(Zmin, MaxDpth), self.elements.z[in_ocean] - VertSwimmingSpeed * abs(self.time_step.total_seconds()))
            
            else:
                # Going up
                in_ocean = np.where(self.elements.z<0)[0] 
                if len(in_ocean) > 0:
                    self.elements.z[in_ocean] = np.minimum(-2.0, self.elements.z[in_ocean] + VertSwimmingSpeed * abs(self.time_step.total_seconds()))


        # New migration based on Sun altitude
        elif MigrationType == 2: 
            # get_position is from suncalc package - returns altitude and azimuth of the Sun
            # print(self.time)
            self.elements.sun_altitude = get_position(self.time, self.elements.lon, self.elements.lat)["altitude"]
            # Vertical motion down durig the day, up at night
            self.elements.vertSwimmingSpeed = - np.sign(self.elements.sun_altitude) * VertSwimmingSpeed 
                  
            # New particle depths
            newZ = self.elements.z + self.elements.vertSwimmingSpeed * abs(self.time_step.total_seconds())
            # Set limits to vertical position
            self.elements.z = np.minimum(-2.0, newZ)  # at least 2 m below surface
            self.elements.z = np.maximum(np.maximum(Zmin, MaxDpth), self.elements.z) # at least 5 m above bottom and max 500 m deep
            print('Sun synchronized migrations ... [max depth: '+str(min(self.elements.z))+']')
        else:
            sys.exit('pelagia_trans.py: Wrong migration type')
               
    def update_pelagia_state(self):
        # Growth, reproduction, mortality 
        # Food dependent. PP is the proxy for food availability.
        
        # Get maximum primary production value in the water column.
        # This will be used for growth / mortality 
        self.elements.p_production = self.environment_profiles['net_primary_production_of_biomass_expressed_as_carbon_per_unit_volume_in_sea_water'].max(axis=0)
              
        
        # Get the coefficients of the growth-mortality equation
        k1 = self.get_config('k1')
        k2 = self.get_config('k2')
        k3 = self.get_config('k3')
           
        # Growth - Michaelis-Menten aka Monod kinetics with additional mortality
        dN = self.elements.biomass * self.time_step.total_seconds()/(3600*24) \
            *  (k1 * self.elements.p_production / (k2 + self.elements.p_production) - k3)
            
        print('k: '+str(k1)+'   k2: '+str(k2)+'   k3: '+str(k3))
                
        # Additional mortality - penalty for shallow water
        Shallow = np.zeros_like(self.environment.sea_floor_depth_below_sea_level)
        Shallow[self.environment.sea_floor_depth_below_sea_level < 200] = -0.03  # daily mortality in shallow water - increased predation; -0.05 halved in 14 days
        dNShallow = self.elements.biomass * self.time_step.total_seconds()/(3600*24) * Shallow
        # Additional mortality - penalty for proximity of the coast - beaching
        NearCoast = self.environment.nearshore_mask * (-0.05)  # daily mortality for proximity of coast
        dNBeaching = self.elements.biomass * self.time_step.total_seconds()/(3600*24) * NearCoast
        
        # record mortality due to beaching and shallow waters
        self.elements.ks = Shallow
        self.elements.kb = NearCoast
        
        # Update biomass (biomass) if each particle / superindividual
        self.elements.biomass = self.elements.biomass + dN + dNShallow + dNBeaching
                
    
    def update(self):
        # Growth and mortality
        self.update_pelagia_state()
        # Advection with horizontal ocean currents
        self.advect_ocean_current()
        self.vertical_mixing()
        self.diel_vertical_migration()