#!/usr/bin/env python
"""
First test run of 
pelagia_trans model
====================
"""
import sys
import numpy as np
from datetime import timedelta, datetime
from opendrift.readers import reader_netCDF_CF_generic
# from opendrift.readers import reader_global_landmask
from opendrift.models.pelagia_IBM import Pelagia_IBM
import opendrift
import random


inFile = sys.argv[1]
startyear = sys.argv[2]
k1 = float(sys.argv[3])
k2 = float(sys.argv[4])
k3 = float(sys.argv[5])


print('Year = '+startyear)
print('k1 = '+str(k1))
print('k2 = '+str(k2))
print('k3 = '+str(k3))

# Get particle locations from previous simulation
odm = Pelagia_IBM(loglevel=20)
odm = opendrift.open(inFile)
elements_final = odm.elements
time_final = datetime( odm.time.year, odm.time.month, odm.time.day )
del odm

# startyear = str(int(time_final.year))
print('Start year: '+startyear)

odm = Pelagia_IBM(loglevel=20) # Initialize model and set log level (0 for debug)

# read MFS files
reader_MFS1 = reader_netCDF_CF_generic.Reader('/media/martin/MP4T_2/data/MFS/cur_rean_daily/'+startyear+'*.nc')
reader_MedBFM1 = reader_netCDF_CF_generic.Reader('/media/martin/MP4T_2/data/MFS-BFM/monthly/'+startyear+'/*BIOL*.nc')
reader_MFS_bath = reader_netCDF_CF_generic.Reader('/media/martin/MP4T_2/data/MFS/statics/MED-MFC_006_004_mask_bathy_copy3.nc')
#

odm.add_reader([reader_MFS1])
odm.add_reader([reader_MedBFM1], variables=['net_primary_production_of_biomass_expressed_as_carbon_per_unit_volume_in_sea_water'])
odm.add_reader(reader_MFS_bath, variables=['sea_floor_depth_below_sea_level'])
odm.add_reader(reader_MFS_bath, variables=['land_binary_mask'])
odm.add_reader(reader_MFS_bath, variables=['nearshore_mask'])

# odm.set_config('general:use_auto_landmask', True) # use default GSHHG vector landmask
odm.set_config('general:use_auto_landmask', False) # use model landmask

# Prepare initial state of particles
lons = np.ma.getdata(elements_final.lon)  # masked array to ordinary array
lats = np.ma.getdata(elements_final.lat)
zs = np.ma.getdata(elements_final.z)
biomasss = np.ma.getdata(elements_final.biomass)

# Remove particles with decimated number of individuals
mask = biomasss > 20
print(mask)
biomasss = biomasss[mask]
lons = lons[mask]
lats = lats[mask]
zs = zs[mask]

removed = biomasss.size - np.sum(mask)
    
print('Removed '+str(removed)+' particles')

# Multiply elements with large number of individuals
print('Adding new particles')
Alons = []
Alats = []
Azs = []
Abiomasss = []
for n in range(len(biomasss)):
    # If there is too many individuals in superindividual. We create additional
    # superindividuals with 10000 individuals
    while biomasss[n] > 20000:
        Alats.append(lats[n] + random.randrange(-250, 250)/1000)
        Alons.append(lons[n] + random.randrange(-250, 250)/1000)
        Azs.append(zs[n]+random.randrange(1,6)-3)
        Abiomasss.append(10000)
        biomasss[n] = biomasss[n] - 10000

# Add new elements to the old array
lons = np.append(lons,Alons)
lats = np.append(lats,Alats)
zs = np.append(zs,Azs)
biomasss = np.append(biomasss,Abiomasss)

# If the experiment is blowing up in numbers, quit!
if len(biomasss) > 200000:
    sys.exit('The population has blown up! Requested # of particles: '+str(len(biomasss)))

# If the population colapsed, quit
if len(biomasss) < 50:
    sys.exit('The population collapsed! Requested # of particles: '+str(len(biomasss)))

print('Seeding '+str(len(biomasss))+' particles')
# Seed particles
odm.seed_elements(lons, lats, z=zs, biomass=biomasss, radius=0, number=len(biomasss), time=datetime(int(startyear),1,1,0,0,0))

# Set beaching behaviour
odm.set_config('general:coastline_action', 'previous')
# exit()

# Set growth-mortality coefficients
odm.set_config('k1', k1)
odm.set_config('k2', k2)
odm.set_config('k3', k3)


print('Starting the simulation ...')
print('Run start: '+str(datetime(int(startyear),1,1,0,0,0)))
print('Run end: '+str(datetime(int(startyear)+1,1,1,0,0,0)))

# Running model (until end of driver data)
odm.run(time_step=timedelta(hours=1), time_step_output=timedelta(days=1), end_time=datetime(int(startyear)+1,1,1,0,0,0), outfile='Pelagia_Med_fwd_'+startyear+'.nc',)


# Print and plot results
print(odm)
#odm.plot(filename='Pelagia_Med_fwd_'+startyear+'.png',linecolor='biomass',lvmin=100,lvmax=20000)
#odm.plot(filename='Pelagia_Med_fwd_BW_'+startyear+'.png')
# odm.plot(linecolor='number')
# odm.animation(fps=30,skip=1,color='number')
# #odm.animation(background=['x_sea_water_velocity','y_sea_water_velocity'], skip=5)
#odm.plot()
