# Encoding: utf-8

import sys
import os
import numpy as np
from scipy import interpolate
import matplotlib.pyplot as plt

from cage.landscape import LandscapeAnalyzer
from cage.core import Facet

"""
Plot the landscape of connected edges. This is some damn illegible code, but
improving it would require thinking, and ain't nobody got time for that.
"""

# PARAMETERS TODO DIRTY AS FUCK -> Find a better way
CATION='Na'
START_FACET = 0 # 0 or 1 -> determines facet to start chain from
ANGLE_MESH = 0.003
RADIUS_MESH = 0.01
CONTOUR_LEVELS = np.mgrid[0.1:2:0.1]
E_LIMITS = (0, 1.4)

# Input

# TODO Make a proper parser
if len(sys.argv) == 3:
    CATION = sys.argv[1]
else:
    CATION = 'Li'

dirname = sys.argv[-1]
dir_list = [os.path.abspath(dir) for dir in os.listdir(dirname)
            if os.path.isdir(dir)]

# TODO Make this part ignore directories that do not have the right files
edges = [[LandscapeAnalyzer.from_file(os.path.join(dir, 'landscape.json')),
          [Facet.from_file(os.path.join(dir, 'init_facet.json')),
          Facet.from_file(os.path.join(dir, 'final_facet.json'))]]
         for dir in dir_list]

for edge in edges:
    edge[0].analyze_cation_energies(facet=edge[1][0],cation=CATION)

# TODO Make landscapes and paths connected via dictionary

# Find all the facets in the paths
facets = []
for edge in edges:
    for facet in edge[1]:
        if facet not in facets:
            facets.append(facet)

# Find end_facets
end_facets = []
for edge in edges:
    for facet in edge[1]:
        if facet not in end_facets:
            end_facets.append(facet)
        else:
            end_facets.remove(facet)

# Check if the chain is connected
if len(end_facets) != 2:
    raise ValueError('Edges are not connected in a chain. Aborting...')
# TODO Handle case of circular connection of paths

# Find the starting path
facet_chain = [end_facets[START_FACET], ]
path_chain = []
landscape_chain = []

while len(facet_chain) < len(facets):
    last_facet = facet_chain[-1]
    print('Last Facet:')
    print(str(last_facet))
    # Find the path that connects to the last facet and is not already in the
    # path chain
    for edge in edges:
        if edge[1] not in path_chain and last_facet in edge[1]:
            other_facet = edge[1].copy()
            other_facet.remove(last_facet)
            other_facet = other_facet[0]
            print('Other Facet:')
            print(str(other_facet))
            facet_chain.append(other_facet)

            if edge[1][0] != last_facet:
                print('Flipped path (before):')
                print(str(edge[1][0]))
                print(str(edge[1][1]))
                edge[1].reverse()
                print('Flipped path (after):')
                print(str(edge[1][0]))
                print(str(edge[1][1]))
                edge[0].flip_coordinates('Angle')

            landscape_chain.append(edge[0])
            path_chain.append(edge[1])

    print('')
print('----------------')
print('Facet Chain:')
for facet in facet_chain:
    print(str(facet))
print('----------------')
print('Path Chain:')
for edge in path_chain:
    print('From')
    print(str(edge[0]))
    print('To')
    print(str(edge[1]))

# Interpolate the landscapes to a uniform mesh

# Find the proper radii
min_max_radius = 1e6
max_min_radius = 0
for landscape in landscape_chain:
    rmax = landscape.datapoints['Distance'].max()
    if rmax < min_max_radius:
        min_max_radius = rmax
    rmin = landscape.datapoints['Distance'].min()
    if rmin > max_min_radius:
        max_min_radius = rmin

print('-----------')
print('Largest minimal radius = ' + str(max_min_radius))
print('Smallest maximal radius = ' + str(min_max_radius))

# Adjust the angles to make one angle coordinate for all edges
facet_angles = [0, landscape_chain[0].datapoints['Angle'].max()]

for landscape in landscape_chain[1:]:
    print('Maximum angle = ' + str(facet_angles[-1]))
    landscape.datapoints['Angle'] += facet_angles[-1]
    facet_angles.append(landscape.datapoints['Angle'].max())

all_radii = []
all_angles = []
all_energy = []

# Interpolate the landscapes
for landscape in landscape_chain:

    data = landscape.datapoints

    data['Distance'] = np.round(data['Distance'], 5)
    data = np.sort(data, order=['Distance', 'Angle'])

    # Find the number of radii and angles
    r_init = data['Distance'][0]
    nangles = 1
    while abs(data['Distance'][nangles] - r_init) < 1e-5:
        nangles += 1
    nradii = int(len(data) / nangles)
    print('')
    print('-----------')
    print('Number of Angles = ' + str(nangles))
    print('Number of Radii = ' + str(nradii))

    # Get the right format for the data
    radii = data['Distance'].reshape(nradii, nangles)  # [::nradii]
    angles = data['Angle'].reshape(nradii, nangles)  # [:nangles]
    energy = data['Energy'].reshape(nradii, nangles)

    print('Shape angles: ' + str(angles.shape))
    print('Shape radii: ' + str(radii.shape))
    print('Shape energy: ' + str(energy.shape))

    new_angles, new_radii = np.mgrid[ angles.min():angles.max():ANGLE_MESH,
                                    max_min_radius:min_max_radius:RADIUS_MESH ]
    print('-------------')
    print('Shape new_angles: ' + str(new_angles.shape))
    print('Shape new_radii: ' + str(new_radii.shape))
    tck = interpolate.bisplrep(angles, radii, energy, s=0.01)

    new_energy = interpolate.bisplev(new_angles[:,0], new_radii[0,:], tck)

    all_radii.append(new_radii)
    all_angles.append(new_angles)
    all_energy.append(new_energy)

total_radii = np.concatenate(tuple(all_radii))
total_angles = np.concatenate(tuple(all_angles))
total_energy = np.concatenate(tuple(all_energy))
total_energy -= total_energy.min()

plt.figure()
plt.pcolor(total_angles, total_radii, total_energy, vmin=E_LIMITS[0],
           vmax=E_LIMITS[1], cmap='viridis')
cbar = plt.colorbar()
cbar.set_label('Energy (eV)', size='x-large')
CS = plt.contour(total_angles, total_radii, total_energy, colors='black',
            levels=CONTOUR_LEVELS, linewidths=0.6)
for angle in facet_angles:
    plt.plot([angle, angle], [total_radii.min(), total_radii.max()], color='k',
         linestyle='-', linewidth=1)
xlabel = []
for i in range(len(facet_angles)):
    xlabel.append('$\Omega_' + str(i+1) + '$')
#plt.xlabel('Angle', size='large')
plt.ylabel('$r$ ($\mathrm{\AA}$)', size='x-large', fontname='Georgia')
plt.xticks(facet_angles, xlabel, size='x-large')
plt.clabel(CS, fontsize=10, inline_spacing=15, fmt='%1.1f', manual=True)
plt.show()


import ctypes
import ctypes.util
from numpy.ctypeslib import ndpointer
import numpy
from scipy import integrate
from .. import potential
from ..potential.planarPotential import planarPotentialFromFullPotential, \
    planarPotentialFromRZPotential
from ..potential.planarPotential import _evaluateplanarRforces,\
    _evaluateplanarphiforces, _evaluateplanarPotentials
from ..potential.WrapperPotential import parentWrapperPotential
from ..util.multi import parallel_map
from ..util.leung_dop853 import dop853
from ..util import symplecticode
from ..util import _load_extension_libs

_lib, _ext_loaded= _load_extension_libs.load_libgalpy()

def _parse_pot(pot):
    """Parse the potential so it can be fed to C"""
    from .integrateFullOrbit import _parse_scf_pot
    #Figure out what's in pot
    if not isinstance(pot,list):
        pot= [pot]
    #Initialize everything
    pot_type= []
    pot_args= []
    npot= len(pot)
    for p in pot:
        # Prepare for wrappers
        if ((isinstance(p,planarPotentialFromFullPotential) \
          or isinstance(p,planarPotentialFromRZPotential)) \
          and isinstance(p._Pot,parentWrapperPotential)) \
        or isinstance(p,parentWrapperPotential):
            if not isinstance(p,parentWrapperPotential):
                wrap_npot, wrap_pot_type, wrap_pot_args= \
                    _parse_pot(potential.toPlanarPotential(p._Pot._pot))
            else:
                wrap_npot, wrap_pot_type, wrap_pot_args= _parse_pot(p._pot)
        if (isinstance(p,planarPotentialFromRZPotential)
            or isinstance(p,planarPotentialFromFullPotential) ) \
                 and isinstance(p._Pot,potential.LogarithmicHaloPotential):
            pot_type.append(0)
            if p._Pot.isNonAxi:
                pot_args.extend([p._Pot._amp,p._Pot._q,
                                 p._Pot._core2,p._Pot._1m1overb2])
            else:
                pot_args.extend([p._Pot._amp,p._Pot._q,p._Pot._core2,2.]) # 1m1overb2 > 1: axi
        elif isinstance(p,planarPotentialFromFullPotential) \
                 and isinstance(p._Pot,potential.DehnenBarPotential):
            pot_type.append(1)
            pot_args.extend([p._Pot._amp*p._Pot._af,p._Pot._tform,
                             p._Pot._tsteady,p._Pot._rb,p._Pot._omegab,
                             p._Pot._barphi])
        elif isinstance(p,potential.TransientLogSpiralPotential):
            pot_type.append(2)
            pot_args.extend([p._amp,p._A,p._to,p._sigma2,p._alpha,p._m,
                             p._omegas,p._gamma])
        elif isinstance(p,potential.SteadyLogSpiralPotential):
            pot_type.append(3)
            if p._tform is None:
                pot_args.extend([p._amp,float('nan'), float('nan'),
                                 p._A,p._alpha,p._m,
                                 p._omegas,p._gamma])
            else:
                pot_args.extend([p._amp,p._tform,p._tsteady,p._A,p._alpha,p._m,
                                 p._omegas,p._gamma])
        elif isinstance(p,potential.EllipticalDiskPotential):
            pot_type.append(4)
            if p._tform is None:
                pot_args.extend([p._amp,float('nan'), float('nan'),
                                 p._twophio,p._p,p._phib])
            else:
                pot_args.extend([p._amp,p._tform,p._tsteady,
                                 p._twophio,p._p,p._phib])
        elif isinstance(p,planarPotentialFromRZPotential) \
                 and isinstance(p._Pot,potential.MiyamotoNagaiPotential):
            pot_type.append(5)
            pot_args.extend([p._Pot._amp,p._Pot._a,p._Pot._b])
        elif isinstance(p,potential.LopsidedDiskPotential):
            pot_type.append(6)
            pot_args.extend([p._amp,p._mphio,p._p,p._phib])
        elif isinstance(p,planarPotentialFromRZPotential) \
                 and isinstance(p._Pot,potential.PowerSphericalPotential):
            pot_type.append(7)
            pot_args.extend([p._Pot._amp,p._Pot.alpha])
        elif isinstance(p,planarPotentialFromRZPotential) \
                 and isinstance(p._Pot,potential.HernquistPotential):
            pot_type.append(8)
            pot_args.extend([p._Pot._amp,p._Pot.a])
        elif isinstance(p,planarPotentialFromRZPotential) \
                 and isinstance(p._Pot,potential.NFWPotential):
            pot_type.append(9)
            pot_args.extend([p._Pot._amp,p._Pot.a])
        elif isinstance(p,planarPotentialFromRZPotential) \
                 and isinstance(p._Pot,potential.JaffePotential):
            pot_type.append(10)
            pot_args.extend([p._Pot._amp,p._Pot.a])
        elif isinstance(p,planarPotentialFromRZPotential) \
                and isinstance(p._Pot,potential.DoubleExponentialDiskPotential):
            pot_type.append(11)
            pot_args.extend([p._Pot._amp,
                             -4.*numpy.pi*p._Pot._alpha*p._Pot._amp,
                             p._Pot._alpha,p._Pot._beta,len(p._Pot._de_j1_xs)])
            pot_args.extend(p._Pot._de_j0_xs)
            pot_args.extend(p._Pot._de_j1_xs)
            pot_args.extend(p._Pot._de_j0_weights)
            pot_args.extend(p._Pot._de_j1_weights)
        elif isinstance(p,planarPotentialFromRZPotential) \
                and isinstance(p._Pot,potential.FlattenedPowerPotential):
            pot_type.append(12)
            pot_args.extend([p._Pot._amp,p._Pot.alpha,p._Pot.core2])
        elif isinstance(p,planarPotentialFromRZPotential) \
                 and isinstance(p._Pot,potential.IsochronePotential):
            pot_type.append(14)
            pot_args.extend([p._Pot._amp,p._Pot.b])
        elif isinstance(p,planarPotentialFromRZPotential) \
                 and isinstance(p._Pot,potential.PowerSphericalPotentialwCutoff):
            pot_type.append(15)
            pot_args.extend([p._Pot._amp,p._Pot.alpha,p._Pot.rc])
        elif isinstance(p,planarPotentialFromRZPotential) \
                 and isinstance(p._Pot,potential.MN3ExponentialDiskPotential):
            # Three Miyamoto-Nagai disks
            npot+= 2
            pot_type.extend([5,5,5])
            pot_args.extend([p._Pot._amp*p._Pot._mn3[0]._amp,
                             p._Pot._mn3[0]._a,p._Pot._mn3[0]._b,
                             p._Pot._amp*p._Pot._mn3[1]._amp,
                             p._Pot._mn3[1]._a,p._Pot._mn3[1]._b,
                             p._Pot._amp*p._Pot._mn3[2]._amp,
                             p._Pot._mn3[2]._a,p._Pot._mn3[2]._b])
        elif isinstance(p,planarPotentialFromRZPotential) \
                 and isinstance(p._Pot,potential.KuzminKutuzovStaeckelPotential):
            pot_type.append(16)
            pot_args.extend([p._Pot._amp,p._Pot._ac,p._Pot._Delta])
        elif isinstance(p,planarPotentialFromRZPotential) \
                 and isinstance(p._Pot,potential.PlummerPotential):
            pot_type.append(17)
            pot_args.extend([p._Pot._amp,p._Pot._b])
        elif isinstance(p,planarPotentialFromRZPotential) \
                 and isinstance(p._Pot,potential.PseudoIsothermalPotential):
            pot_type.append(18)
            pot_args.extend([p._Pot._amp,p._Pot._a])
        elif isinstance(p,planarPotentialFromRZPotential) \
                 and isinstance(p._Pot,potential.KuzminDiskPotential):
            pot_type.append(19)
            pot_args.extend([p._Pot._amp,p._Pot._a])
        elif isinstance(p,planarPotentialFromRZPotential) \
                 and isinstance(p._Pot,potential.BurkertPotential):
            pot_type.append(20)
            pot_args.extend([p._Pot._amp,p._Pot.a])
        elif (isinstance(p,planarPotentialFromFullPotential) or isinstance(p,planarPotentialFromRZPotential)) \
                and isinstance(p._Pot,potential.EllipsoidalPotential.EllipsoidalPotential):
            pot_args.append(p._Pot._amp)
            pot_args.extend([0.,0.,0.,0.,0.,0.]) # for caching
            if isinstance(p._Pot,potential.TriaxialHernquistPotential):
                pot_type.append(21)
                pot_args.extend([2,p._Pot.a,p._Pot.a4]) # for psi, mdens, mdens_deriv
            if isinstance(p._Pot,potential.TriaxialNFWPotential):
                pot_type.append(22)
                pot_args.extend([2,p._Pot.a,p._Pot.a3]) # for psi, mdens, mdens_deriv
            if isinstance(p._Pot,potential.TriaxialJaffePotential):
                pot_type.append(23)
                pot_args.extend([2,p._Pot.a,p._Pot.a2]) # for psi, mdens, mdens_deriv
            elif isinstance(p._Pot,potential.PerfectEllipsoidPotential):
                pot_type.append(30)
                pot_args.extend([1,p._Pot.a2]) # for psi, mdens, mdens_deriv
            elif isinstance(p._Pot,potential.TriaxialGaussianPotential):
                pot_type.append(37)
                pot_args.extend([1,-p._Pot._twosigma2]) # for psi, mdens, mdens_deriv
            pot_args.extend([p._Pot._b2,p._Pot._c2,
                             int(p._Pot._aligned)]) # Reg. Ellipsoidal
            if not p._Pot._aligned:
                pot_args.extend(list(p._Pot._rot.flatten()))
            else:
                pot_args.extend(list(numpy.eye(3).flatten())) # not actually used
            pot_args.append(p._Pot._glorder)
            pot_args.extend([p._Pot._glx[ii] for ii in range(p._Pot._glorder)])
            # this adds some common factors to the integration weights
            pot_args.extend([-4.*numpy.pi*p._Pot._glw[ii]*p._Pot._b*p._Pot._c\
                            /numpy.sqrt(( 1.+(p._Pot._b2-1.)*p._Pot._glx[ii]**2.)
                                     *(1.+(p._Pot._c2-1.)*p._Pot._glx[ii]**2.))
                             for ii in range(p._Pot._glorder)])
        elif (isinstance(p,planarPotentialFromFullPotential) or isinstance(p,planarPotentialFromRZPotential)) \
                 and isinstance(p._Pot,potential.SCFPotential):
            pt,pa= _parse_scf_pot(p._Pot)
            pot_type.append(pt)
            pot_args.extend(pa)
        elif isinstance(p,planarPotentialFromFullPotential) \
                 and isinstance(p._Pot,potential.SoftenedNeedleBarPotential):
            pot_type.append(25)
            pot_args.extend([p._Pot._amp,p._Pot._a,p._Pot._b,p._Pot._c2,
                             p._Pot._pa,p._Pot._omegab])
            pot_args.extend([0.,0.,0.,0.,0.,0.,0.]) # for caching
        elif (isinstance(p,planarPotentialFromFullPotential) or isinstance(p,planarPotentialFromRZPotential)) \
                and isinstance(p._Pot,potential.DiskSCFPotential):
            # Need to pull this apart into: (a) SCF part, (b) constituent
            # [Sigma_i,h_i] parts
            # (a) SCF, multiply in any add'l amp
            pt,pa= _parse_scf_pot(p._Pot._scf,extra_amp=p._Pot._amp)
            pot_type.append(pt)
            pot_args.extend(pa)
            # (b) constituent [Sigma_i,h_i] parts
            for Sigma,hz in zip(p._Pot._Sigma_dict,p._Pot._hz_dict):
                npot+= 1
                pot_type.append(26)
                stype= Sigma.get('type','exp')
                if stype == 'exp' and not 'Rhole' in Sigma:
                    pot_args.extend([3,0,
                                     4.*numpy.pi*Sigma.get('amp',1.)*p._Pot._amp,
                                     Sigma.get('h',1./3.)])
                elif stype == 'expwhole' \
                        or (stype == 'exp' and 'Rhole' in Sigma):
                    pot_args.extend([4,1,
                                     4.*numpy.pi*Sigma.get('amp',1.)*p._Pot._amp,
                                     Sigma.get('h',1./3.),
                                     Sigma.get('Rhole',0.5)])
                hztype= hz.get('type','exp')
                if hztype == 'exp':
                    pot_args.extend([0,hz.get('h',0.0375)])
                elif hztype == 'sech2':
                    pot_args.extend([1,hz.get('h',0.0375)])
        elif isinstance(p,planarPotentialFromFullPotential) \
                and isinstance(p._Pot, potential.SpiralArmsPotential):
            pot_type.append(27)
            pot_args.extend([len(p._Pot._Cs), p._Pot._amp, p._Pot._N, p._Pot._sin_alpha,
                             p._Pot._tan_alpha, p._Pot._r_ref, p._Pot._phi_ref, p._Pot._Rs, p._Pot._H, p._Pot._omega])
            pot_args.extend(p._Pot._Cs)
        elif isinstance(p,potential.CosmphiDiskPotential):
            pot_type.append(28)
            pot_args.extend([p._amp,p._mphio,p._p,p._mphib,p._m,
                             p._rb,p._rbp,p._rb2p,p._r1p])
        elif isinstance(p,potential.HenonHeilesPotential):
            pot_type.append(29)
            pot_args.extend([p._amp])
        # 30: PerfectEllipsoidPotential, done with other EllipsoidalPotentials above
        # 31: KGPotential
        # 32: IsothermalDiskPotential
        elif isinstance(p, planarPotentialFromRZPotential) \
                and isinstance(p._Pot,potential.DehnenCoreSphericalPotential):
            pot_type.append(33)
            pot_args.extend([p._Pot._amp,p._Pot.a])
        elif isinstance(p, planarPotentialFromRZPotential) \
                and isinstance(p._Pot,potential.DehnenSphericalPotential):
            pot_type.append(34)
            pot_args.extend([p._Pot._amp,p._Pot.a,p._Pot.alpha])
        # 35: HomogeneousSpherePotential
        elif isinstance(p,planarPotentialFromRZPotential) \
             and isinstance(p._Pot,potential.HomogeneousSpherePotential):
            pot_type.append(35)
            pot_args.extend([p._Pot._amp,p._Pot._R2,p._Pot._R3])
        # 36: interpSphericalPotential
        elif isinstance(p,planarPotentialFromRZPotential) \
             and isinstance(p._Pot,potential.interpSphericalPotential):
            pot_type.append(36)
            pot_args.append(len(p._Pot._rgrid))
            pot_args.extend(p._Pot._rgrid)
            pot_args.extend(p._Pot._rforce_grid)
            pot_args.extend([p._Pot._amp,p._Pot._rmin,p._Pot._rmax,
                             p._Pot._total_mass,p._Pot._Phi0,p._Pot._Phimax])
        # 37: TriaxialGaussianPotential, done with other EllipsoidalPotentials above
        ############################## WRAPPERS ###############################
        elif ((isinstance(p,planarPotentialFromFullPotential) or isinstance(p,planarPotentialFromRZPotential)) \
              and isinstance(p._Pot,potential.DehnenSmoothWrapperPotential)) \
              or isinstance(p,potential.DehnenSmoothWrapperPotential):
            if not isinstance(p,potential.DehnenSmoothWrapperPotential):
                p= p._Pot
            pot_type.append(-1)
            # wrap_pot_type, args, and npot obtained before this horrible if
            pot_args.append(wrap_npot)
            pot_type.extend(wrap_pot_type)
            pot_args.extend(wrap_pot_args)
            pot_args.extend([p._amp,p._tform,p._tsteady,int(p._grow)])
        elif ((isinstance(p,planarPotentialFromFullPotential) or isinstance(p,planarPotentialFromRZPotential)) \
          and isinstance(p._Pot,potential.SolidBodyRotationWrapperPotential)) \
          or isinstance(p,potential.SolidBodyRotationWrapperPotential):
            if not isinstance(p,potential.SolidBodyRotationWrapperPotential):
                p= p._Pot
            pot_type.append(-2)
            # wrap_pot_type, args, and npot obtained before this horrible if
            pot_args.append(wrap_npot)
            pot_type.extend(wrap_pot_type)
            pot_args.extend(wrap_pot_args)
            pot_args.extend([p._amp,p._omega,p._pa])
        elif ((isinstance(p,planarPotentialFromFullPotential) or isinstance(p,planarPotentialFromRZPotential)) \
          and isinstance(p._Pot,potential.CorotatingRotationWrapperPotential)) \
          or isinstance(p,potential.CorotatingRotationWrapperPotential):
            if not isinstance(p,potential.CorotatingRotationWrapperPotential):
                p= p._Pot
            pot_type.append(-4)
            # wrap_pot_type, args, and npot obtained before this horrible if
            pot_args.append(wrap_npot)
            pot_type.extend(wrap_pot_type)
            pot_args.extend(wrap_pot_args)
            pot_args.extend([p._amp,p._vpo,p._beta,p._pa,p._to])
        elif ((isinstance(p,planarPotentialFromFullPotential) or isinstance(p,planarPotentialFromRZPotential)) \
              and isinstance(p._Pot,potential.GaussianAmplitudeWrapperPotential)) \
              or isinstance(p,potential.GaussianAmplitudeWrapperPotential):
            if not isinstance(p,potential.GaussianAmplitudeWrapperPotential):
                p= p._Pot
            pot_type.append(-5)
            # wrap_pot_type, args, and npot obtained before this horrible if
            pot_args.append(wrap_npot)
            pot_type.extend(wrap_pot_type)
            pot_args.extend(wrap_pot_args)
            pot_args.extend([p._amp,p._to,p._sigma2])
        elif ((isinstance(p,planarPotentialFromFullPotential) or isinstance(p,planarPotentialFromRZPotential)) \
              and isinstance(p._Pot,potential.MovingObjectPotential)) \
              or isinstance(p,potential.MovingObjectPotential):
            if not isinstance(p,potential.MovingObjectPotential):
                p= p._Pot
            pot_type.append(-6)
            wrap_npot, wrap_pot_type, wrap_pot_args= \
                    _parse_pot(potential.toPlanarPotential(p._pot))
            pot_args.append(wrap_npot)
            pot_type.extend(wrap_pot_type)
            pot_args.extend(wrap_pot_args)
            pot_args.extend([len(p._orb.t)])
            pot_args.extend(p._orb.t)
            pot_args.extend(p._orb.x(p._orb.t,use_physical=False))
            pot_args.extend(p._orb.y(p._orb.t,use_physical=False))
            pot_args.extend([p._amp])
            pot_args.extend([p._orb.t[0],p._orb.t[-1]]) #t_0, t_f
    pot_type= numpy.array(pot_type,dtype=numpy.int32,order='C')
    pot_args= numpy.array(pot_args,dtype=numpy.float64,order='C')
    return (npot,pot_type,pot_args)

def _parse_integrator(int_method):
    """parse the integrator method to pass to C"""
    #Pick integrator
    if int_method.lower() == 'rk4_c':
        int_method_c= 1
    elif int_method.lower() == 'rk6_c':
        int_method_c= 2
    elif int_method.lower() == 'symplec4_c':
        int_method_c= 3
    elif int_method.lower() == 'symplec6_c':
        int_method_c= 4
    elif int_method.lower() == 'dopr54_c':
        int_method_c= 5
    elif int_method.lower() == 'dop853_c':
        int_method_c= 6
    else:
        int_method_c= 0
    return int_method_c

def _parse_tol(rtol,atol):
    """Parse the tolerance keywords"""
    #Process atol and rtol
    if rtol is None:
        rtol= -12.*numpy.log(10.)
    else: #pragma: no cover
        rtol= numpy.log(rtol)
    if atol is None:
        atol= -12.*numpy.log(10.)
    else: #pragma: no cover
        atol= numpy.log(atol)
    return (rtol,atol)

def integratePlanarOrbit_c(pot,yo,t,int_method,rtol=None,atol=None,
                           dt=None):
    """
    NAME:
       integratePlanarOrbit_c
    PURPOSE:
       C integrate an ode for a planarOrbit
    INPUT:
       pot - Potential or list of such instances
       yo - initial condition [q,p], can be [N,4] or [4]
       t - set of times at which one wants the result
       int_method= 'leapfrog_c', 'rk4_c', 'rk6_c', 'symplec4_c', ...
       rtol, atol
       dt= (None) force integrator to use this stepsize (default is to automatically determine one)
   OUTPUT:
       (y,err)
       y : array, shape (len(y0),len(t),4)
       Array containing the value of y for each desired time in t, \
       with the initial value y0 in the first row.
       err: error message, if not zero: 1 means maximum step reduction happened for adaptive integrators
    HISTORY:
       2011-10-03 - Written - Bovy (IAS)
       2018-12-20 - Adapted to allow multiple objects - Bovy (UofT)
    """
    if len(yo.shape) == 1: single_obj= True
    else: single_obj= False
    yo= numpy.atleast_2d(yo)
    nobj= len(yo)
    rtol, atol= _parse_tol(rtol,atol)
    npot, pot_type, pot_args= _parse_pot(pot)
    int_method_c= _parse_integrator(int_method)
    if dt is None:
        dt= -9999.99

    #Set up result array
    result= numpy.empty((nobj,len(t),4))
    err= numpy.zeros(nobj,dtype=numpy.int32)

    #Set up the C code
    ndarrayFlags= ('C_CONTIGUOUS','WRITEABLE')
    integrationFunc= _lib.integratePlanarOrbit
    integrationFunc.argtypes= [ctypes.c_int,
                               ndpointer(dtype=numpy.float64,flags=ndarrayFlags),
                               ctypes.c_int,
                               ndpointer(dtype=numpy.float64,flags=ndarrayFlags),
                               ctypes.c_int,
                               ndpointer(dtype=numpy.int32,flags=ndarrayFlags),
                               ndpointer(dtype=numpy.float64,flags=ndarrayFlags),
                               ctypes.c_double,
                               ctypes.c_double,
                               ctypes.c_double,
                               ndpointer(dtype=numpy.float64,flags=ndarrayFlags),
                               ndpointer(dtype=numpy.int32,flags=ndarrayFlags),
                               ctypes.c_int]

    #Array requirements, first store old order
    f_cont= [yo.flags['F_CONTIGUOUS'],
             t.flags['F_CONTIGUOUS']]
    yo= numpy.require(yo,dtype=numpy.float64,requirements=['C','W'])
    t= numpy.require(t,dtype=numpy.float64,requirements=['C','W'])
    result= numpy.require(result,dtype=numpy.float64,requirements=['C','W'])
    err= numpy.require(err,dtype=numpy.int32,requirements=['C','W'])

    #Run the C code
    integrationFunc(ctypes.c_int(nobj),
                    yo,
                    ctypes.c_int(len(t)),
                    t,
                    ctypes.c_int(npot),
                    pot_type,
                    pot_args,
                    ctypes.c_double(dt),
                    ctypes.c_double(rtol),
                    ctypes.c_double(atol),
                    result,
                    err,
                    ctypes.c_int(int_method_c))

    if numpy.any(err == -10): #pragma: no cover
        raise KeyboardInterrupt("Orbit integration interrupted by CTRL-C (SIGINT)")

    #Reset input arrays
    if f_cont[0]: yo= numpy.asfortranarray(yo)
    if f_cont[1]: t= numpy.asfortranarray(t)

    if single_obj: return (result[0],err[0])
    else: return (result,err)

def integratePlanarOrbit_dxdv_c(pot,yo,dyo,t,int_method,rtol=None,atol=None,
                                dt=None):
    """
    NAME:
       integratePlanarOrbit_dxdv_c
    PURPOSE:
       C integrate an ode for a planarOrbit+phase space volume dxdv
    INPUT:
       pot - Potential or list of such instances
       yo - initial condition [q,p]
       dyo - initial condition [dq,dp]
       t - set of times at which one wants the result
       int_method= 'leapfrog_c', 'rk4_c', 'rk6_c', 'symplec4_c'
       rtol, atol
       dt= (None) force integrator to use this stepsize (default is to automatically determine one))
    OUTPUT:
       (y,err)
       y,dy : array, shape (len(y0),len(t),8)
       Array containing the value of y for each desired time in t, \
       with the initial value y0 in the first row.
       err: error message if not zero, 1: maximum step reduction happened for adaptive integrators
    HISTORY:
       2011-10-19 - Written - Bovy (IAS)
    """
    rtol, atol= _parse_tol(rtol,atol)
    npot, pot_type, pot_args= _parse_pot(pot)
    int_method_c= _parse_integrator(int_method)
    if dt is None:
        dt= -9999.99
    yo= numpy.concatenate((yo,dyo))

    #Set up result array
    result= numpy.empty((len(t),8))
    err= ctypes.c_int(0)

    #Set up the C code
    ndarrayFlags= ('C_CONTIGUOUS','WRITEABLE')
    integrationFunc= _lib.integratePlanarOrbit_dxdv
    integrationFunc.argtypes= [ndpointer(dtype=numpy.float64,flags=ndarrayFlags),
                               ctypes.c_int,
                               ndpointer(dtype=numpy.float64,flags=ndarrayFlags),
                               ctypes.c_int,
                               ndpointer(dtype=numpy.int32,flags=ndarrayFlags),
                               ndpointer(dtype=numpy.float64,flags=ndarrayFlags),
                               ctypes.c_double,
                               ctypes.c_double,
                               ctypes.c_double,
                               ndpointer(dtype=numpy.float64,flags=ndarrayFlags),
                               ctypes.POINTER(ctypes.c_int),
                               ctypes.c_int]

    #Array requirements, first store old order
    f_cont= [yo.flags['F_CONTIGUOUS'],
             t.flags['F_CONTIGUOUS']]
    yo= numpy.require(yo,dtype=numpy.float64,requirements=['C','W'])
    t= numpy.require(t,dtype=numpy.float64,requirements=['C','W'])
    result= numpy.require(result,dtype=numpy.float64,requirements=['C','W'])

    #Run the C code
    integrationFunc(yo,
                    ctypes.c_int(len(t)),
                    t,
                    ctypes.c_int(npot),
                    pot_type,
                    pot_args,
                    ctypes.c_double(dt),
                    ctypes.c_double(rtol),ctypes.c_double(atol),
                    result,
                    ctypes.byref(err),
                    ctypes.c_int(int_method_c))

    if err.value == -10: #pragma: no cover
        raise KeyboardInterrupt("Orbit integration interrupted by CTRL-C (SIGINT)")

    #Reset input arrays
    if f_cont[0]: yo= numpy.asfortranarray(yo)
    if f_cont[1]: t= numpy.asfortranarray(t)

    return (result,err.value)

def integratePlanarOrbit(pot,yo,t,int_method,rtol=None,atol=None,numcores=1,
                         dt=None):
    """
    NAME:
       integratePlanarOrbit
    PURPOSE:
       Integrate an ode for a planarOrbit
    INPUT:
       pot - Potential or list of such instances
       yo - initial condition [q,p], shape [N,3] or [N,4]
       t - set of times at which one wants the result
       int_method= 'leapfrog', 'odeint', or 'dop853'
       rtol, atol= tolerances (not always used...)
       numcores= (1) number of cores to use for multi-processing
       dt= (None) force integrator to use this stepsize (default is to automatically determine one; only for C-based integrators!)
    OUTPUT:
       (y,err)
       y : array, shape (N,len(t),3/4)
       Array containing the value of y for each desired time in t, \
       with the initial value y0 in the first row.
       err: error message, always zero for now
    HISTORY:
       2010-07-20 - Written - Bovy (NYU)
       2019-04-09 - Adapted to allow multiple objects and parallel mapping - Bovy (UofT)
    """
    nophi= False
    if not int_method.lower() == 'dop853' and not int_method == 'odeint':
        if len(yo[0]) == 3:
            nophi= True
            #We hack this by putting in a dummy phi=0
            yo= numpy.pad(yo,((0,0),(0,1)),'constant',constant_values=0)
    if int_method.lower() == 'leapfrog':
        if rtol is None: rtol= 1e-8
        def integrate_for_map(vxvv):
            #go to the rectangular frame
            this_vxvv= numpy.array([vxvv[0]*numpy.cos(vxvv[3]),
                                 vxvv[0]*numpy.sin(vxvv[3]),
                                 vxvv[1]*numpy.cos(vxvv[3])
                                     -vxvv[2]*numpy.sin(vxvv[3]),
                                 vxvv[2]*numpy.cos(vxvv[3])
                                     +vxvv[1]*numpy.sin(vxvv[3])])
            #integrate
            tmp_out= symplecticode.leapfrog(_planarRectForce,this_vxvv,
                                            t,args=(pot,),rtol=rtol)
            #go back to the cylindrical frame
            R= numpy.sqrt(tmp_out[:,0]**2.+tmp_out[:,1]**2.)
            phi= numpy.arccos(tmp_out[:,0]/R)
            phi[(tmp_out[:,1] < 0.)]= 2.*numpy.pi-phi[(tmp_out[:,1] < 0.)]
            vR= tmp_out[:,2]*numpy.cos(phi)+tmp_out[:,3]*numpy.sin(phi)
            vT= tmp_out[:,3]*numpy.cos(phi)-tmp_out[:,2]*numpy.sin(phi)
            out= numpy.zeros((len(t),4))
            out[:,0]= R
            out[:,1]= vR
            out[:,2]= vT
            out[:,3]= phi
            return out
    elif int_method.lower() == 'dop853' or int_method.lower() == 'odeint':
        if rtol is None: rtol= 1e-8
        if int_method.lower() == 'dop853':
            integrator= dop853
            extra_kwargs= {}
        else:
            integrator= integrate.odeint
            extra_kwargs= {'rtol':rtol}
        if len(yo[0]) == 3:
            def integrate_for_map(vxvv):
                l= vxvv[0]*vxvv[2]
                l2= l**2.
                init= [vxvv[0],vxvv[1]]
                intOut= integrator(_planarREOM,init,t=t,args=(pot,l2),
                                   **extra_kwargs)
                out= numpy.zeros((len(t),3))
                out[:,0]= intOut[:,0]
                out[:,1]= intOut[:,1]
                out[:,2]= l/out[:,0]
                #post-process to remove negative radii
                neg_radii= (out[:,0] < 0.)
                out[neg_radii,0]= -out[neg_radii,0]
                return out
        else:
            def integrate_for_map(vxvv):
                vphi= vxvv[2]/vxvv[0]
                init= [vxvv[0],vxvv[1],vxvv[3],vphi]
                intOut= integrator(_planarEOM,init,t=t,args=(pot,),
                                   **extra_kwargs)
                out= numpy.zeros((len(t),4))
                out[:,0]= intOut[:,0]
                out[:,1]= intOut[:,1]
                out[:,3]= intOut[:,2]
                out[:,2]= out[:,0]*intOut[:,3]
                #post-process to remove negative radii
                neg_radii= (out[:,0] < 0.)
                out[neg_radii,0]= -out[neg_radii,0]
                out[neg_radii,3]+= numpy.pi
                return out
    else: # Assume we are forcing parallel_mapping of a C integrator...
        def integrate_for_map(vxvv):
            return integratePlanarOrbit_c(pot,numpy.copy(vxvv),
                                          t,int_method,dt=dt)[0]
    if len(yo) == 1: # Can't map a single value...
        out= numpy.atleast_3d(integrate_for_map(yo[0]).T).T
    else:
        out= numpy.array((parallel_map(integrate_for_map,yo,numcores=numcores)))
    if nophi:
        out= out[:,:,:3]
    return out, numpy.zeros(len(yo))

def integratePlanarOrbit_dxdv(pot,yo,dyo,t,int_method,
                              rectIn,rectOut,
                              rtol=None,atol=None,
                              dt=None,numcores=1):
    """
    NAME:
       integratePlanarOrbit_dxdv
    PURPOSE:
       Integrate an ode for a planarOrbit+phase space volume dxdv
    INPUT:
       pot - Potential or list of such instances
       yo - initial condition [q,p], shape [N,4]
       dyo - initial condition [dq,dp], shape [N,4]
       t - set of times at which one wants the result
       int_method= 'odeint', 'dop853', 'dopr54_c', 'rk4_c', 'rk6_c'
       rectIn= (False) if True, input dyo is in rectangular coordinates
       rectOut= (False) if True, output dyo is in rectangular coordinates
       rtol, atol= tolerances (not always used...)
       numcores= (1) number of cores to use for multi-processing
       dt= (None) force integrator to use this stepsize (default is to automatically determine one; only for C-based integrators)
    OUTPUT:
       (y,err)
       y : array, shape (N,len(t),8)
       Array containing the value of y for each desired time in t, \
       with the initial value y0 in the first row.
       err: error message, always zero for now
    HISTORY:
       2011-10-17 - Written - Bovy (IAS)
       2019-05-21 - Adapted to allow multiple objects and parallel mapping - Bovy (UofT)
    """
    #go to the rectangular frame
    this_yo= numpy.array([yo[:,0]*numpy.cos(yo[:,3]),
                         yo[:,0]*numpy.sin(yo[:,3]),
                         yo[:,1]*numpy.cos(yo[:,3])
                           -yo[:,2]*numpy.sin(yo[:,3]),
                         yo[:,2]*numpy.cos(yo[:,3])
                           +yo[:,1]*numpy.sin(yo[:,3])]).T
    if not rectIn:
        this_dyo= numpy.array([numpy.cos(yo[:,3])*dyo[:,0]
                              -yo[:,0]*numpy.sin(yo[:,3])*dyo[:,3],
                            numpy.sin(yo[:,3])*dyo[:,0]
                              +yo[:,0]*numpy.cos(yo[:,3])*dyo[:,3],
                            -(yo[:,1]*numpy.sin(yo[:,3])
                              +yo[:,2]*numpy.cos(yo[:,3]))*dyo[:,3]
                              +numpy.cos(yo[:,3])*dyo[:,1]
                              -numpy.sin(yo[:,3])*dyo[:,2],
                            (yo[:,1]*numpy.cos(yo[:,3])
                              -yo[:,2]*numpy.sin(yo[:,3]))*dyo[:,3]
                              +numpy.sin(yo[:,3])*dyo[:,1]
                              +numpy.cos(yo[:,3])*dyo[:,2]]).T
    else:
        this_dyo= dyo
    this_yo= numpy.hstack((this_yo,this_dyo))
    if int_method.lower() == 'dop853' or int_method.lower() == 'odeint':
        if rtol is None: rtol= 1e-8
        if int_method.lower() == 'dop853':
            integrator= dop853
            extra_kwargs= {}
        else:
            integrator= integrate.odeint
            extra_kwargs= {'rtol':rtol}
        def integrate_for_map(vxvv):
            return integrator(_planarEOM_dxdv,vxvv,t=t,args=(pot,),
                              **extra_kwargs)
    else: # Assume we are forcing parallel_mapping of a C integrator...
        def integrate_for_map(vxvv):
            return integratePlanarOrbit_dxdv_c(pot,numpy.copy(vxvv[:4]),
                                               numpy.copy(vxvv[4:]),
                                               t,int_method,dt=dt,
                                               rtol=rtol,atol=atol)[0]
    if len(this_yo) == 1: # Can't map a single value...
        out= numpy.atleast_3d(integrate_for_map(this_yo[0]).T).T
    else:
        out= numpy.array((parallel_map(integrate_for_map,this_yo,
                                    numcores=numcores)))
    #go back to the cylindrical frame
    R= numpy.sqrt(out[...,0]**2.+out[...,1]**2.)
    phi= numpy.arccos(out[...,0]/R)
    phi[(out[...,1] < 0.)]= 2.*numpy.pi-phi[(out[...,1] < 0.)]
    vR= out[...,2]*numpy.cos(phi)+out[...,3]*numpy.sin(phi)
    vT= out[...,3]*numpy.cos(phi)-out[...,2]*numpy.sin(phi)
    cp= numpy.cos(phi)
    sp= numpy.sin(phi)
    out[...,0]= R
    out[...,1]= vR
    out[...,2]= vT
    out[...,3]= phi
    if rectOut:
        out[...,4:]= out[...,4:]
    else:
        dR= cp*out[...,4]+sp*out[...,5]
        dphi= (cp*out[...,5]-sp*out[...,4])/R
        dvR= cp*out[...,6]+sp*out[...,7]+vT*dphi
        dvT= cp*out[...,7]-sp*out[...,6]-vR*dphi
        out[...,4]= dR
        out[...,7]= dphi
        out[...,5]= dvR
        out[...,6]= dvT
    return out, numpy.zeros(len(yo))

def _planarREOM(y,t,pot,l2):
    """
    NAME:
       _planarREOM
    PURPOSE:
       implements the EOM, i.e., the right-hand side of the differential
       equation, for integrating a planar Orbit assuming angular momentum
       conservation
    INPUT:
       y - current phase-space position
       t - current time
       pot - (list of) Potential instance(s)
       l2 - angular momentum squared
    OUTPUT:
       dy/dt
    HISTORY:
       2010-07-20 - Written - Bovy (NYU)
    """
    return [y[1],
            l2/y[0]**3.+_evaluateplanarRforces(pot,y[0],t=t)]

def _planarEOM(y,t,pot):
    """
    NAME:
       _planarEOM
    PURPOSE:
       implements the EOM, i.e., the right-hand side of the differential
       equation, for integrating a general planar Orbit
    INPUT:
       y - current phase-space position
       t - current time
       pot - (list of) Potential instance(s)
    OUTPUT:
       dy/dt
    HISTORY:
       2010-07-20 - Written - Bovy (NYU)
    """
    l2= (y[0]**2.*y[3])**2.
    return [y[1],
            l2/y[0]**3.+_evaluateplanarRforces(pot,y[0],phi=y[2],t=t),
            y[3],
            1./y[0]**2.*(_evaluateplanarphiforces(pot,y[0],phi=y[2],t=t)-
                         2.*y[0]*y[1]*y[3])]

def _planarEOM_dxdv(x,t,pot):
    """
    NAME:
       _planarEOM_dxdv
    PURPOSE:
       implements the EOM, i.e., the right-hand side of the differential
       equation, for integrating phase space differences, rectangular
    INPUT:
       x - current phase-space position
       t - current time
       pot - (list of) Potential instance(s)
    OUTPUT:
       dy/dt
    HISTORY:
       2011-10-18 - Written - Bovy (IAS)
    """
    #x is rectangular so calculate R and phi
    R= numpy.sqrt(x[0]**2.+x[1]**2.)
    phi= numpy.arccos(x[0]/R)
    sinphi= x[1]/R
    cosphi= x[0]/R
    if x[1] < 0.: phi= 2.*numpy.pi-phi
    #calculate forces
    Rforce= _evaluateplanarRforces(pot,R,phi=phi,t=t)
    phiforce= _evaluateplanarphiforces(pot,R,phi=phi,t=t)
    R2deriv= _evaluateplanarPotentials(pot,R,phi=phi,t=t,dR=2)
    phi2deriv= _evaluateplanarPotentials(pot,R,phi=phi,t=t,dphi=2)
    Rphideriv= _evaluateplanarPotentials(pot,R,phi=phi,t=t,dR=1,dphi=1)
    #Calculate derivatives and derivatives+time derivatives
    dFxdx= -cosphi**2.*R2deriv\
           +2.*cosphi*sinphi/R**2.*phiforce\
           +sinphi**2./R*Rforce\
           +2.*sinphi*cosphi/R*Rphideriv\
           -sinphi**2./R**2.*phi2deriv
    dFxdy= -sinphi*cosphi*R2deriv\
           +(sinphi**2.-cosphi**2.)/R**2.*phiforce\
           -cosphi*sinphi/R*Rforce\
           -(cosphi**2.-sinphi**2.)/R*Rphideriv\
           +cosphi*sinphi/R**2.*phi2deriv
    dFydx= -cosphi*sinphi*R2deriv\
           +(sinphi**2.-cosphi**2.)/R**2.*phiforce\
           +(sinphi**2.-cosphi**2.)/R*Rphideriv\
           -sinphi*cosphi/R*Rforce\
           +sinphi*cosphi/R**2.*phi2deriv
    dFydy= -sinphi**2.*R2deriv\
           -2.*sinphi*cosphi/R**2.*phiforce\
           -2.*sinphi*cosphi/R*Rphideriv\
           +cosphi**2./R*Rforce\
           -cosphi**2./R**2.*phi2deriv
    return numpy.array([x[2],x[3],
                     cosphi*Rforce-1./R*sinphi*phiforce,
                     sinphi*Rforce+1./R*cosphi*phiforce,
                     x[6],x[7],
                     dFxdx*x[4]+dFxdy*x[5],
                     dFydx*x[4]+dFydy*x[5]])

def _planarRectForce(x,pot,t=0.):
    """
    NAME:
       _planarRectForce
    PURPOSE:
       returns the planar force in the rectangular frame
    INPUT:
       x - current position
       t - current time
       pot - (list of) Potential instance(s)
    OUTPUT:
       force
    HISTORY:
       2011-02-02 - Written - Bovy (NYU)
    """
    #x is rectangular so calculate R and phi
    R= numpy.sqrt(x[0]**2.+x[1]**2.)
    phi= numpy.arccos(x[0]/R)
    sinphi= x[1]/R
    cosphi= x[0]/R
    if x[1] < 0.: phi= 2.*numpy.pi-phi
    #calculate forces
    Rforce= _evaluateplanarRforces(pot,R,phi=phi,t=t)
    phiforce= _evaluateplanarphiforces(pot,R,phi=phi,t=t)
    return numpy.array([cosphi*Rforce-1./R*sinphi*phiforce,
                     sinphi*Rforce+1./R*cosphi*phiforce])


"""This module contains iceberg drift models.
"""

import cmath
import numpy as np
from icedef.constants import *

def newtonian_drift_wrapper(t, lon, lat, vx, vy, **kwargs):
    """This function performs interpolations for current and wind velocities and then runs the drift model.

    Args:
        t (numpy.datetime64): time.
        lon (float): longitude.
        lat (float): latitude.
        vx (float): x-component of iceberg velocity in m/s.
        vy (float): y-component of iceberg velocity in m/s.
        **kwargs: coming soon - see source code for now.

    Returns:
        vx (float): new x-component of iceberg velocity in m/s.
        vy (float): new y-component of iceberg velocity in m/s.
        ax (float): x-component of iceberg acceleration in m/s.
        ay (float): y-component of iceberg acceleration in m/s.
    """

    dt = kwargs.pop('time_step', np.timedelta64(300, 's'))

    fast_interpolation = kwargs.pop('fast_interpolation', True)

    if fast_interpolation:

        current_interpolator = kwargs.pop('current_interpolator')
        Vcx, Vcy = current_interpolator((t, lat, lon))
        wind_interpolator = kwargs.pop('wind_interpolator')
        Vwx, Vwy = wind_interpolator((t, lat, lon))

        Vcx_left, Vcy_left = current_interpolator((t - dt, lat, lon))
        Vcx_right, Vcy_right = current_interpolator((t + dt, lat, lon))

    else:

        Vcxs = kwargs.pop('eastward_current')
        Vcys = kwargs.pop('northward_current')
        Vwxs = kwargs.pop('eastward_wind')
        Vwys = kwargs.pop('northward_wind')

        Vcx = Vcxs.interp(time=t, latitude=lat, longitude=lon, assume_sorted=True).values
        Vcy = Vcys.interp(time=t, latitude=lat, longitude=lon, assume_sorted=True).values
        Vwx = Vwxs.interp(time=t, latitude=lat, longitude=lon, assume_sorted=True).values
        Vwy = Vwys.interp(time=t, latitude=lat, longitude=lon, assume_sorted=True).values

        Vcx_left = Vcxs.interp(time=t - dt, latitude=lat, longitude=lon, assume_sorted=True).values
        Vcx_right = Vcxs.interp(time=t + dt, latitude=lat, longitude=lon, assume_sorted=True).values
        Vcy_left = Vcys.interp(time=t - dt, latitude=lat, longitude=lon, assume_sorted=True).values
        Vcy_right = Vcys.interp(time=t + dt, latitude=lat, longitude=lon, assume_sorted=True).values

    Amwx = (Vcx_right - Vcx_left) / (dt.item().total_seconds() * 2)
    Amwy = (Vcy_right - Vcy_left) / (dt.item().total_seconds() * 2)

    kwargs['Vcx'] = Vcx
    kwargs['Vcy'] = Vcy
    kwargs['Vwx'] = Vwx
    kwargs['Vwy'] = Vwy
    kwargs['Amwx'] = Amwx
    kwargs['Amwy'] = Amwy

    kwargs['phi'] = lat

    ax, ay = newtonian_drift(vx, vy, **kwargs)

    return vx, vy, ax, ay

def newtonian_drift(Vx, Vy, **kwargs):
    """This function computes iceberg acceleration using a general Newtonian drift model.

    Args:
        Vx (float): x-component of iceberg velocity in m/s.
        Vy (float): y-component of iceberg velocity in m/s.
        **kwargs: coming soon - see source code for now.

    Returns:
        ax (float): x-component of iceberg acceleration in m/s.
        ay (float): y-component of iceberg acceleration in m/s.
    """

    # Constants
    Omega = EARTH_ROTATION_RATE
    rhoa = AIR_DENSITY
    rhow = SEAWATER_DENSITY

    Vwx = kwargs.pop('Vwx')
    Vwy = kwargs.pop('Vwy')
    Vcx = kwargs.pop('Vcx')
    Vcy = kwargs.pop('Vcy')

    Amwx = kwargs.pop('Amwx', 0)
    Amwy = kwargs.pop('Amwy', 0)

    Ca = kwargs.pop('form_drag_coefficient_in_air', 1.5)
    Cw = kwargs.pop('form_drag_coefficient_in_water', 1.5)
    Cda = kwargs.pop('skin_drag_coefficient_in_air', 2.5e-4)
    Cdw = kwargs.pop('skin_drag_coefficient_in_water', 5e-4)
    As = kwargs.pop('sail_area', 9600.0)
    Ak = kwargs.pop('keel_area', 48000.0)
    At = kwargs.pop('top_area', 25600.0)
    Ab = kwargs.pop('bottom_area', 25600.0)
    M = kwargs.pop('mass', 5468160000.0)
    phi = kwargs.pop('latitude', 50)

    # Wind force
    Fax = (0.5 * rhoa * Ca * As + rhoa * Cda * At) * np.sqrt((Vwx - Vx)**2 + (Vwy - Vy)**2) * (Vwx - Vx)
    Fay = (0.5 * rhoa * Ca * As + rhoa * Cda * At) * np.sqrt((Vwx - Vx)**2 + (Vwy - Vy)**2) * (Vwy - Vy)

    # Current force
    ekman = kwargs.pop('ekman', False)

    if ekman:

        Fwx_list = []
        Fwy_list = []
        Vcx_list = []
        Vcy_list = []

        depth_vec = kwargs.pop('depth_vec', np.arange(0, -110, -10))

        u_vec, v_vec = compute_ekman_spiral((Vwx, Vwy), (Vcx, Vcy), depth_vec)

        for i in range(len(u_vec)):

            Vcx, Vcy = u_vec[i], v_vec[i]

            Vcx_list.append(Vcx)
            Vcy_list.append(Vcy)

            Fwx = (0.5 * rhow * Cw * Ak + rhow * Cdw * Ab) * np.sqrt((Vcx - Vx)**2 + (Vcy - Vy)**2) * (Vcx - Vx)
            Fwy = (0.5 * rhow * Cw * Ak + rhow * Cdw * Ab) * np.sqrt((Vcx - Vx)**2 + (Vcy - Vy)**2) * (Vcy - Vy)

            Fwx_list.append(Fwx)
            Fwy_list.append(Fwy)

        Vcx = np.mean(np.array(Vcx_list))
        Vcy = np.mean(np.array(Vcy_list))

        Fwx = np.mean(np.array(Fwx_list))
        Fwy = np.mean(np.array(Fwy_list))

    else:
        Fwx = (0.5 * rhow * Cw * Ak + rhow * Cdw * Ab) * np.sqrt((Vcx - Vx)**2 + (Vcy - Vy)**2) * (Vcx - Vx)
        Fwy = (0.5 * rhow * Cw * Ak + rhow * Cdw * Ab) * np.sqrt((Vcx - Vx)**2 + (Vcy - Vy)**2) * (Vcy - Vy)

    # Coriolis Parameter
    f = 2 * Omega * np.sin(np.deg2rad(phi))

    # Coriolis force
    Fcx = f * M * Vy
    Fcy = -f * M * Vx

    # Water Pressure Gradient Force
    Vmwx = Vcx
    Vmwy = Vcy
    Amwx = Amwx
    Amwy = Amwy
    Fwpx = M * (Amwx - f * Vmwy)
    Fwpy = M * (Amwy + f * Vmwx)

    # Iceberg acceleration
    ax = (Fax + Fwx + Fcx + Fwpx) / (M + 0.5 * M)
    ay = (Fay + Fwy + Fcy + Fwpy) / (M + 0.5 * M)

    log = kwargs.pop('log', None)

    if log is not None:
        log.info('{:.2f},{:.2f},{:.2f},{:.2f},{:.2f},{:.2f},{:.2f},{:.2f}'.format(
            Fax, Fay, Fwx, Fwy, Fcx, Fcy, Fwpx, Fwpy))

    return ax, ay

def compute_ekman_velocity(wind, depth, latitude=50):
    """This function computes Ekman velocity at some depth.

    Args:
        wind (tuple of float): components (x, y) of wind velocity in m/s.
        depth (float): depth below the sea surface (down is negative) in m.
        latitude: latitude.

    Returns:
        u_ekman (float): x-component of Ekman velocity at specified depth.
        v_ekman (float): y-component of Ekman velocity at specified depth.
    """

    u_wind, v_wind = wind
    z = depth
    phi = latitude

    rho_air = 1.225  # density of air (kg/m^3)
    rho_water = 1028  # density of seawater (kg/m^3)
    Cd = 1.3e-3  # ranges from (1.1 - 1.5) x 10^-3
    Omega = 7.2910e-5  # rotation rate of Earth (s^-1)
    Az = 5e-2  # m^2 s^−1;

    f = lambda phi: 2 * Omega * np.sin(phi)  # Coriolis parameter

    tau_x = lambda U, V: rho_air * Cd * U * np.sqrt(U ** 2 + V ** 2)  # wind stress x-component
    tau_y = lambda U, V: rho_air * Cd * V * np.sqrt(U ** 2 + V ** 2)  # wind stress y-component

    V0x = tau_x(u_wind, v_wind) / np.sqrt(rho_water ** 2 * np.abs(f(phi)) * Az)
    V0y = tau_y(u_wind, v_wind) / np.sqrt(rho_water ** 2 * np.abs(f(phi)) * Az)
    V0 = np.sqrt(V0x ** 2 + V0y ** 2)
    theta = np.pi / 2 - np.arctan2(V0y, V0x)

    a = np.sqrt(abs(f(phi)) / (2 * Az))

    # note: clockwise rotation
    u_ekman = V0 * np.exp(a * z) * np.cos(np.pi / 4 + a * z) * (np.cos(theta) + np.sin(theta))
    v_ekman = V0 * np.exp(a * z) * np.sin(np.pi / 4 + a * z) * (np.cos(theta) - np.sin(theta))

    return u_ekman, v_ekman

def compute_ekman_spiral(wind, surface_current, depth_vec, latitude=50):
    """This function computes an Ekman spiral.

    Args:
        wind (tuple of float): components (x, y) of wind velocity in m/s.
        surface_current (tuple of float): components (x, y) of current velocity at the surface in m/s.
        depth_vec (list of float): depths in m to compute Ekman velocity at.
        latitude: latitude.

    Returns:
        u_current_vec (list of float): x-components of Ekman velocity at all depths specified.
        v_current_vec (list of float): y-components of Ekman velocity at all depths specified.
    """

    u_ekman_vec = np.zeros(len(depth_vec))
    v_ekman_vec = np.zeros(len(depth_vec))

    u_ekman_surface, v_ekman_surface = compute_ekman_velocity(wind, 0, latitude)

    u_surface_current, v_surface_current = surface_current

    u_barotropic = u_surface_current - u_ekman_surface
    v_barotropic = v_surface_current - v_ekman_surface

    for i, depth in enumerate(depth_vec):
        u_ekman_vec[i], v_ekman_vec[i] = compute_ekman_velocity(wind, depth, latitude)

    u_current_vec = u_barotropic + u_ekman_vec
    v_current_vec = v_barotropic + v_ekman_vec

    return u_current_vec, v_current_vec

def analytical_drift_wrapper(t, lon, lat, **kwargs):
    """This function performs interpolations for current and wind velocities and then runs the drift model.

    Args:
        t (numpy.datetime64): time.
        lon (float): longitude.
        lat (float): latitude.
        vx (float): x-component of iceberg velocity in m/s.
        vy (float): y-component of iceberg velocity in m/s.
        **kwargs: coming soon - see source code for now.

    Returns:
        vx (float): new x-component of iceberg velocity in m/s.
        vy (float): new y-component of iceberg velocity in m/s.

    """

    fast_interpolation = kwargs.pop('fast_interpolation', True)

    if fast_interpolation:

        current_interpolator = kwargs.pop('current_interpolator')
        wind_interpolator = kwargs.pop('wind_interpolator')

        vwu, vwv = current_interpolator((t, lat, lon))
        vau, vav = wind_interpolator((t, lat, lon))

    else:

        vwus = kwargs.pop('eastward_current')
        vwvs = kwargs.pop('northward_current')
        vaus = kwargs.pop('eastward_wind')
        vavs = kwargs.pop('northward_wind')

        vwu = vwus.interp(time=t, latitude=lat, longitude=lon, assume_sorted=True).values
        vwv = vwvs.interp(time=t, latitude=lat, longitude=lon, assume_sorted=True).values
        vau = vaus.interp(time=t, latitude=lat, longitude=lon, assume_sorted=True).values
        vav = vavs.interp(time=t, latitude=lat, longitude=lon, assume_sorted=True).values

    kwargs['vwu'] = vwu
    kwargs['vwv'] = vwv
    kwargs['vau'] = vau
    kwargs['vav'] = vav

    vx, vy = analytical_drift(lon, lat, **kwargs)

    return vx, vy

def analytical_drift(x, y, **kwargs):
    """This function computes the velocity of an iceberg using an analytical drift model.

    Args:
        x (float): longitude
        y (float): latitude
        **kwargs: coming soon - see source code for now.

    Returns:
        viu (float): x-component of iceberg velocity in m/s.
        viv (float): y-component of iceberg velocity in m/s.
    """

    vwu = kwargs.pop('vwu')
    vwv = kwargs.pop('vwv')
    vau = kwargs.pop('vau')
    vav = kwargs.pop('vav')

    l = kwargs.pop('waterline_length', 160)
    w = kwargs.pop('waterline_length', 160)  # note: currently all icebergs are cuboid

    Cw = kwargs.pop('form_drag_coefficient_in_water', 0.9)
    Ca = kwargs.pop('form_drag_coefficient_in_air', 1.3)

    Omega = EARTH_ROTATION_RATE
    rhow = SEAWATER_DENSITY
    rhoa = AIR_DENSITY
    rhoi = ICEBERG_DENSITY

    gamma = np.sqrt(rhoa * (rhow - rhoi) / rhow / rhoi * (Ca / Cw))
    S = np.pi * ((l * w) / (l + w))
    f = 2 * Omega * np.sin((np.abs(y) * np.pi) / 180)
    Lambda = np.sqrt(2) * Cw * (gamma * np.sqrt(vau ** 2 + vav ** 2)) / (f * S)

    if Lambda < 0.1:
        alpha = Lambda * (Lambda**4 * (Lambda**4 * (Lambda**4 * (-0.0386699020961393 * Lambda**4 +
                0.055242717280199) - 0.0883883476483184) + 0.176776695296637) - 0.707106781186548)

    else:
        alpha = np.multiply(np.divide(np.sqrt(2), np.power(Lambda, 3)), (1 - np.sqrt(1 + np.power(Lambda, 4))))

    if Lambda < 0.6:
        beta = Lambda**3 * (Lambda**4 * (Lambda**4 * (Lambda**4 * (Lambda**4 *
                (Lambda**4 * (Lambda**4 * (Lambda**4 * (Lambda**4 * (0.0153268598203613 *
                Lambda**4 - 0.0151656272365985) + 0.0180267866272764) + 0.0219176256311202) -
                0.0274446790511418) + 0.0357675015202851) - 0.0493731785691779) + 0.0745776683282687) -
                0.132582521472478) + 0.353553390593274)

    else:
        beta = np.real(np.multiply(np.divide(1, np.power(Lambda, 3)), cmath.sqrt(np.multiply((4 +
                np.power(Lambda, 4)), cmath.sqrt(1 + np.power(Lambda, 4))) - 3 * np.power(Lambda, 4) - 4)))

    viu = vwu + gamma * (-alpha * vav + beta * vau)
    viv = vwv + gamma * (alpha * vau + beta * vav)

    return viu, viv


# discern.py
# Contact: Jacob Schreiber
#          jmschreiber91@gmail.com

'''
Run specific analyses for a cancer dataset using the general functions defined
in ogimos. Each function should be a specific cancer type. These will usually
run analyses in batches and output RData files.
'''

import pandas as pd
import numpy as np
import multiprocessing
import rpy2.robjects as ro
import matplotlib.pyplot as plt

from rpy2.robjects.numpy2ri import numpy2ri
ro.conversion.py2ri = numpy2ri

def scale( data ):
	'''
	Scale a matrix in a columnwise manner.
	'''

	return ( data - data.mean( axis=0 ) ) / data.std( axis=0 )

class DISCERN( object ):
	'''
	DISCERN is an unsupervised feature selection algorithm which uses
	differential correlation in order to identify perturbed features given
	data from two conditions. A Gaussian Graphical Model (GGM) is constructed
	for each condition, with nodes being features and edges being weights
	calculated by an Elastic Net Regressor. The DISCERN score is then calculated
	by looking at which edges differ significantly between the two networks, and
	identifying which genes are perturbed between the features, i.e. have
	different neighbors in the GGM.
	'''

	def __init__( self ):
		pass

	def fit_score( self, null_training, null_testing, alternate_training,
		alternate_testing, names, mask=None, l=0, alpha=1.00, n_cores=None ):
		'''
		Build a GGM for the null dataset and the alternate dataset separately
		using an elastic net regressor. This is done through the R package
		glmnet. The DISCERN score will then be calculated and saved in
		self._scores.
		'''

		# Determine the number of cores to use. If -1 is passed in, use all the
		# cores, if nothing is passed in use 1 core, else use the specified
		# number of cores.
		if n_cores == -1:
			n_cores = multiprocessing.cpu_count()
		else:
			n_cores = n_cores or 1

		# If lambda is specified as 'auto', determine the best lambda using the
		# null training set.
		if l in ['auto', 'Auto']:
			l, sse = self.lambda_opt( null_training, names, mask, n_cores,
				alpha=alpha )

		# Assign a uniform true mask on the covariates by default
		mask = mask or np.ones( null_training.shape[1] )

		# First we need to import the R libraries we want to work with
		ro.r( "library(glmnet)" )
		ro.r( "library(foreach)" )
		ro.r( "library(doParallel)" )

		# Set up the cluster using the given size
		ro.r( "cl <- makeCluster({})".format( n_cores ) )
		ro.r( "registerDoParallel(cl)" )

		# Scale all four data sets used independently of each other.
		null_testing = ( null_testing - null_training.mean( axis=0 ) ) / null_training.std( axis=0 )
		null_training = ( null_training - null_training.mean( axis=0 ) ) / null_training.std( axis=0 )
		alternate_testing = ( alternate_testing - alternate_training.mean( axis=0 ) ) / alternate_training.std( axis=0 )
		alternate_training = ( alternate_training - alternate_training.mean( axis=0) ) / alternate_training.std( axis=0 )

		# Now we need to push the data we're working with to the R environment
		ro.r.assign( "null_training", null_training )
		ro.r.assign( "alternate_training", alternate_training )
		ro.r.assign( "null_testing", null_testing )
		ro.r.assign( "alternate_testing", alternate_testing )
		ro.r.assign( "names", names )
		ro.r.assign( "mask", mask )

		# Now push the glmnet usage and scoring function
		# y_n and y_a refer to a specific column (gene) from those two matrices
		# X_n and X_a refer to all the covariates, excluding the one chosen for y
		# fit_n and fit_a are fitten glmnet objects for the null or cancer set
		# y_pred_ab are the predicted values using fit_a on X_b
		# error_ab is the error calculated using fit_a on X_b

		ro.r( r"""
		discern <- function( lambda, alpha ) {
			clusterExport( cl, c("null_testing", "null_training", "alternate_testing",
				"alternate_training", "names", "mask" ))
			results <- foreach( i=1:dim(null_training)[2], .packages='glmnet' ) %dopar% {
				covariates <- mask
				covariates[i] = F
				name = names[[i]]

				y_n <- null_training[, i]
				y_a <- alternate_training[, i]

				X_n <- null_training[, as.logical( covariates ) ]
				X_a <- alternate_training[, as.logical( covariates ) ]

				fit_n <- glmnet( X_n, y_n, standardize=FALSE, alpha=alpha, lambda=lambda )
				fit_a <- glmnet( X_a, y_a, standardize=FALSE, alpha=alpha, lambda=lambda )

				y_n <- null_testing[, i]
				y_a <- alternate_testing[, i]

				X_n <- null_testing[, as.logical( covariates ) ]
				X_a <- alternate_testing[, as.logical( covariates ) ]

				r <- tryCatch(
				{
					y_pred_nn <- predict( fit_n, X_n, s=lambda )
					y_pred_na <- predict( fit_n, X_a, s=lambda )
					y_pred_an <- predict( fit_a, X_n, s=lambda )
					y_pred_aa <- predict( fit_a, X_a, s=lambda )

					error_nn = sum( (y_pred_nn-y_n)^2 )
					error_na = sum( (y_pred_na-y_a)^2 )
					error_an = sum( (y_pred_an-y_n)^2 )
					error_aa = sum( (y_pred_aa-y_a)^2 )

					T4 = ( error_na + error_an ) / ( error_nn + error_aa )
					T2 = error_na + error_an - error_nn - error_aa
					c( T2, T4 )

				}, error = function(err) {
					c( NaN, NaN )
				} )

				result = c( name, r[1], r[2] )
				return( result )
			}

			stopCluster(cl)
			gc()
			return( results )
		}""" )

		scores = np.array( ro.r['discern']( l, alpha  ) )
		self._scores = pd.DataFrame( scores, columns=['Feature', 'T2', 'T4'] )
		self._scores.index = self._scores['Feature']
		self._scores = self._scores.convert_objects(convert_numeric=True)
		return self._scores

	def lambda_opt( self, data, names, mask=None, n_cores=None,
		nfolds=5, alpha=1.00, plot=False ):
		'''
		Determine lambda_opt for a given set of data. Make sure that the data
		used here is the training data for the
		'''

		# Determine the number of cores to use. If -1 is passed in, use all the
		# cores, if nothing is passed in use 1 core, else use the specified
		# number of cores.
		if n_cores == -1:
			n_cores = multiprocessing.cpu_count()
		else:
			n_cores = n_cores or 1

		# Assign a uniform true mask on the covariates by default
		mask = mask or np.ones( data.shape[1] )
		lambdas = 10**np.arange( 1, -3.5, -.05 )

		# First we need to import the R libraries we want to work with
		ro.r( "library(glmnet)" )
		ro.r( "library(foreach)" )
		ro.r( "library(doParallel)" )
		ro.r( "library(matrixStats)")

		# Set up the cluster using the given size
		ro.r( "cl <- makeCluster({})".format( n_cores ) )
		ro.r( "registerDoParallel(cl)" )

		# Pass these variables into R
		ro.r.assign( "data", data )
		ro.r.assign( "names", names )
		ro.r.assign( "mask", mask )
		ro.r.assign( "alpha", alpha )
		ro.r.assign( "lambdas", lambdas )
		ro.r.assign( "nfolds", nfolds )

		lambda_cv = ro.r("""
		lambda_cv <- function() {
			clusterExport( cl, c("data", "names", "mask", "alpha", "lambdas", "nfolds") )

			SSE <- foreach( i=1:dim(data)[2], .packages=c('glmnet', 'foreach', 'matrixStats'), .combine='+' ) %dopar% {
				covariates <- mask
				covariates[i] = F
				name = names[[i]]

				errors <- foreach( j=1:nfolds, .combine='+' ) %do% {
					fold = seq( j, dim(data)[1], nfolds )

					y <- data[ -fold, i]
					X <- data[ -fold, as.logical( covariates ) ]

					x_mu_fit <- colMeans( X )
					x_sigma_fit <- colSds( X )

					y_mu_fit <- mean(y)
					y_sigma_fit <- sd(y)

					y <- ( y - y_mu_fit ) / y_sigma_fit
					X <- t( ( t(X) - x_mu_fit ) / x_sigma_fit )

					fit <- glmnet( X, y, standardize=FALSE, alpha=alpha, lambda=lambdas )

					y <- data[ fold, i ]
					X <- data[ fold, as.logical( covariates ) ]

					y <- ( y - y_mu_fit ) / y_sigma_fit
					X <- t( ( t(X) - x_mu_fit ) / x_sigma_fit )

					y_pred = predict( fit, X )

					e <- foreach( k=1:dim(y_pred)[2] ) %do% {
						sum( ( y_pred[, k] - y )^2 )
					}

					e = unlist(e)
					return( e )
				}

				errors = unlist(errors)
				return( errors )
 			}

			stopCluster(cl)
			return( SSE )
		}
		""")

		l = np.array(ro.r['lambda_cv']())

		if plot:
			plt.plot( lambdas, l, c='c', alpha=0.5 )
			plt.xscale('log')
			plt.xlabel('$\lambda$')
			plt.ylabel('SSE')
			plt.title('Cross Validation Selection of $\lambda$')
			plt.savefig('lambda_opt.png')

		return lambdas[ np.argmin(l) ], l.min()

	def sparsity( self, data, names, mask=None, n_cores=None,
		nfolds=5, alpha=1.00 ):
		'''
		Create a graph of the sparsity of the matrix.
		'''

		# Determine the number of cores to use. If -1 is passed in, use all the
		# cores, if nothing is passed in use 1 core, else use the specified
		# number of cores.
		if n_cores == -1:
			n_cores = multiprocessing.cpu_count()
		else:
			n_cores = n_cores or 1

		# Assign a uniform true mask on the covariates by default
		mask = mask or np.ones( data.shape[1] )
		lambdas = 10**np.arange( 1, -3.5, -.05 )

		# First we need to import the R libraries we want to work with
		ro.r( "library(glmnet)" )
		ro.r( "library(foreach)" )
		ro.r( "library(doParallel)" )
		ro.r( "library(matrixStats)")

		# Set up the cluster using the given size
		ro.r( "cl <- makeCluster({})".format( n_cores ) )
		ro.r( "registerDoParallel(cl)" )

		# Pass these variables into R
		ro.r.assign( "data", data )
		ro.r.assign( "names", names )
		ro.r.assign( "mask", mask )
		ro.r.assign( "alpha", alpha )
		ro.r.assign( "lambdas", lambdas )
		ro.r.assign( "nfolds", nfolds )

		sparsity_cv = ro.r("""
		sparsity_cv <- function() {
			clusterExport( cl, c("data", "names", "mask", "alpha", "lambdas", "nfolds") )

			#sparsity <- foreach( i=1:dim(data)[2], .packages=c('glmnet', 'foreach', 'matrixStats'), .combine='+' ) %dopar% {
			sparsity <- foreach( i=1:dim(data)[2], .packages=c('glmnet', 'foreach', 'matrixStats'), .combine='+' ) %dopar% {
				covariates <- mask
				covariates[i] = F
				name = names[[i]]

				edges <- foreach( j=1:nfolds, .combine='+' ) %do% {
					fold = seq( j, dim(data)[1], nfolds )

					y <- data[ -fold, i]
					X <- data[ -fold, as.logical( covariates ) ]

					x_mu_fit <- colMeans( X )
					x_sigma_fit <- colSds( X )

					y_mu_fit <- mean(y)
					y_sigma_fit <- sd(y)

					y <- ( y - y_mu_fit ) / y_sigma_fit
					X <- t( ( t(X) - x_mu_fit ) / x_sigma_fit )

					fit <- glmnet( X, y, standardize=FALSE, alpha=alpha, lambda=lambdas )

					n_edges = fit['df']$df / nfolds
					return(n_edges)
				}

				edges = unlist(edges)
				return( edges )
 			}

			stopCluster(cl)
			return( sparsity )
		}
		""")

		s = np.array(ro.r['sparsity_cv']())
		s /= data.shape[1] ** 2
		print s

		plt.plot( lambdas, s, c='c', alpha=0.5, linewidth=2.5 )
		plt.xscale('log')
		plt.xlabel('$\lambda$')
		plt.ylabel('Percent Edges')
		plt.savefig('sparsity.png')


"""
TIMP (Time-Inhomogeneous Markov Process) Data_Stall Recovery Trigger Modeling
~~~~~~~~~~~~~~~~~~~~~
TIMP model is for determining proper triggers of Data_Stall recovery

Currently, Android uses 1 minute as the sole trigger for recovery, which is found to be inefficient.
To overcome this, we develop a TIMP a formalize and model the entire Data_Stall-Recovery process.
Based on this, we calculate the most suitable triggers that may result in the global minimum of recovery time.

Note that to run this model, you may need Data_Stall duration data (which we will not disclose yet), as well as
devices' recovery stages

:copyright: (c) 2020 by Cellular Reliability Team.
:license: Apache 2.0, see LICENSE for more details.
"""
import numpy as np
import os
import multiprocessing
from functools import partial

def integrate_over_threshold_for_stage(thresholds, stage):
    """
    Calculate the integration of event duration over different thresholds (triggers) for each recovery stage
    :param thresholds: recovery triggers for entering the next recovery stage
    :param stage: current recovery stage
    :return: integration value
    """
    durations_for_stage = durations_for_stages[stage]
    start = 0
    for i in range(stage):
        start += thresholds[i]
    if stage != 3:
        end = start + thresholds[stage]
    else:
        end = durations_for_stage[-1]

    durations_over_threshold = durations_for_stage[np.where(durations_for_stage <= end)[0]]
    durations_over_threshold = durations_over_threshold[np.where(durations_over_threshold >= start)[0]]
    if len(durations_over_threshold) == 0:
        return 0
    return np.average(durations_over_threshold)

def penalties():
    """
    Penalties for executing each recovery operation
    :return: penalties
    """
    durations_over_threshold = [duration for duration in durations if duration <= 60000]
    penalty = np.average(durations_over_threshold)
    return [0, penalty, 2 * penalty, 3 * penalty]

def cdf_for_stages(thresholds):
    """
    Calculate the CDF of event durations for each recovery stage
    Here recovery stages are divided using given thresholds
    :param thresholds: recovery triggers for entering the next recovery stage
    :return: CDF
    """
    end = 0
    cdf = list()
    for stage in range(3):
        end += thresholds[stage]
        cdf.append(len(np.where(durations <= end)[0]) / len(durations))
    return cdf

def overhead(thresholds, cdf, stage=0):
    """
    The expected recovery time starting at corresponding stage
    This function is in recursive form so that we know the overall expected recovery time when stage is 0
    :param thresholds: recovery triggers for entering the next recovery stage
    :param cdf: CDF for durations in different recovery stages
    :param stage: current recovery stage
    :return: the expected recovery time
    """
    if len(thresholds) != 3 or len(durations_for_stages) != 4:
        raise Exception("Threshold values not enough!")

    if stage == 3:
        return recovery_penalties[stage] + \
               integrate_over_threshold_for_stage(thresholds, stage)

    return recovery_penalties[stage] + \
           integrate_over_threshold_for_stage(thresholds, stage) + \
           (1 - cdf[stage]) * overhead(thresholds, cdf, stage + 1)

def data_processing():
    """
    Dedicated data processing function that turns CSV to NPY file for fast data loading
    :return: None
    """
    for file_name in os.listdir(os.getcwd()):
        if not (file_name.startswith("DATA_STALL") and file_name.endswith(".csv")):
            continue
        file_name = file_name[:-4]
        data = list()
        with open("{}.csv".format(file_name), 'r') as file:
            for line in file:
                line = int(line.strip())
                if line <= 86400000:
                    data.append(line)
        data.reverse()
        np.save("{}-mod.npy".format(file_name), np.array(data))

def prepare_data():
    """
    Prepare all local data
    To run this, you'll need corresponding Data_Stall duration data for each recovery stage
    :return: None
    """
    global durations, recovery_penalties
    files = list()
    for file in os.listdir(os.getcwd()):
        if file.endswith("mod.npy"):
            files.append(file)
    files = np.sort(files)
    for file in files:
        if file.endswith("all-mod.npy"):
            durations = np.load(file)
            continue
        print("Preparing {}...".format(file))
        durations_for_stages.append(np.load(file))

    print("Preparing penalties...")
    recovery_penalties = penalties()
    print("Preparing CDF...")
    print("End data preparations.")

def loss(threshold1, threshold2, threshold3):
    """
    The loss function that should be minimized for different thresholds, i.e., the expected recovery time
    :param threshold1: the first trigger
    :param threshold2: the second trigger
    :param threshold3: the third trigger
    :return: loss
    """
    thresholds = [threshold1, threshold2, threshold3]
    print(thresholds)
    cdf = cdf_for_stages(thresholds)
    result = overhead(thresholds, cdf, 0)
    return result

def brute_force():
    """
    Used for searching the global minimum (can be replaced by dual annealing or others)
    :return: loss values for the entire trigger space
    """
    cores = multiprocessing.cpu_count()
    p = multiprocessing.Pool(processes=cores)
    result = list()
    for i in range(0, 60000, 1000):
        for j in range(0, 60000, 1000):
            par = partial(loss, threshold2=i, threshold3=j)
            res = p.map(par, range(0, 60000, 1000))
            result.extend(res)
    return result

if __name__ == '__main__':
    # overall durations
    durations = list()
    # durations for events recovered in each stage
    durations_for_stages = list()
    recovery_penalties = list()
    prepare_data()
    # currently we use brute force to better present our results (since all search space is considered)
    # we strongly recommend using other global optimization algorithms such as dual annealing for faster iteration
    results = brute_force()
    # print the global minimum and its corresponding index
    print(np.min(results))
    print(np.argmin(results))
    np.save("results.npy", results)


"""
Latent Dirichlet Allocation models, including vanilla LDA as well as correlated
and dynamic topic models (CTMs and DTMs) employing the Polya-Gamma augmentation.
"""

import abc
import copy
import numpy as np
from scipy.misc import logsumexp
from scipy.linalg import solve_triangular as _solve_triangular
from scipy.linalg.lapack import dpotrs as _dpotrs
from scipy.special import gammaln
import scipy.sparse

from pypolyagamma import pgdrawvpar
from gslrandom import multinomial
from pybasicbayes.distributions import Gaussian
from pylds.lds_messages_interface import filter_and_sample_randomwalk

from pgmult.utils import \
    kappa_vec, N_vec, \
    compute_uniform_mean_psi, psi_to_pi, pi_to_psi, \
    ln_pi_to_psi, ln_psi_to_pi, \
    initialize_polya_gamma_samplers, \
    initialize_pyrngs

###
# Util
###

def dpotrs(L, a):
    return _dpotrs(L, a, lower=True)[0]

def solve_triangular(L, a):
    return _solve_triangular(L, a, lower=True, trans='T')

def sample_dirichlet(a, normalize):
    if 'vertical'.startswith(normalize):
        return np.hstack([np.random.dirichlet(col)[:,None] for col in a.T])
    else:
        return np.vstack([np.random.dirichlet(row) for row in a])

def sample_infogaussian(J, h, randvec=None):
    randvec = randvec if randvec is not None else np.random.randn(h.shape[0])
    L = np.linalg.cholesky(J)
    return dpotrs(L, h) + solve_triangular(L, randvec)

def csr_nonzero(mat):
    rows = np.arange(mat.shape[0]).repeat(np.diff(mat.indptr))
    cols = mat.indices
    return rows, cols

def normalize_rows(a):
    a /= a.sum(1)[:,None]
    return a

def log_likelihood(data, wordprobs):
    return np.sum(np.nan_to_num(np.log(wordprobs)) * data.data) \
        + gammaln(data.sum(1)+1).sum() - gammaln(data.data+1).sum()

def check_timestamps(timestamps):
    assert np.all(timestamps == timestamps[np.argsort(timestamps)])

def is_sorted(a):
    return np.all(a == a[np.argsort(a)])

def timeindices_from_timestamps(timestamps):
    return np.array(timestamps) - timestamps[0]

###
# LDA Models
###

# LDA and the CTMs all treat beta, z, and likelihoods the same way, so that
# stuff is factored out into _LDABase. Since each model treats theta
# differently, theta stuff is left abstract.

class _LDABase(object):
    def __init__(self, data, T, alpha_beta):
        assert isinstance(data, scipy.sparse.csr.csr_matrix)
        self.D, self.V = data.shape
        self.T = T
        self.alpha_beta = alpha_beta

        self.data = data

        self.pyrngs = initialize_pyrngs()

        self.initialize_beta()
        self.initialize_theta()
        self.z = np.zeros((data.data.shape[0], T), dtype='uint32')
        self.resample_z()

        # precompute
        self._training_gammalns = \
            gammaln(data.sum(1)+1).sum() - gammaln(data.data+1).sum()

    @abc.abstractproperty
    def theta(self):
        pass

    @abc.abstractmethod
    def initialize_theta(self):
        pass

    @abc.abstractmethod
    def resample_theta(self):
        pass

    def initialize_beta(self):
        self.beta = sample_dirichlet(
            self.alpha_beta * np.ones((self.V, self.T)), 'vert')

    @abc.abstractproperty
    def copy_sample(self):
        pass

    def get_wordprobs(self, data):
        return self.theta.dot(self.beta.T)[csr_nonzero(data)]

    def get_topicprobs(self, data):
        rows, cols = csr_nonzero(data)
        return normalize_rows(self.theta[rows] * self.beta[cols])

    def log_likelihood(self, data=None):
        if data is not None:
            return log_likelihood(data, self.get_wordprobs(data))
        else:
            # this version avoids recomputing the training gammalns
            wordprobs = self.get_wordprobs(self.data)
            return np.sum(np.nan_to_num(np.log(wordprobs)) * self.data.data) \
                + self._training_gammalns

    def perplexity(self, data):
        return np.exp(-self.log_likelihood(data)
                      / data.sum())

    def resample(self):
        self.resample_z()
        self.resample_theta()
        self.resample_beta()

    def resample_beta(self):
        self.beta = sample_dirichlet(
            self.alpha_beta + self.word_topic_counts, 'v')

    def resample_z(self):
        topicprobs = self.get_topicprobs(self.data)
        multinomial(self.pyrngs, self.data.data, topicprobs, self.z)
        self._update_counts()

    def _update_counts(self):
        self.doc_topic_counts = np.zeros((self.D, self.T), dtype='uint32')
        self.word_topic_counts = np.zeros((self.V, self.T), dtype='uint32')
        rows, cols = csr_nonzero(self.data)
        for i, j, zvec in zip(rows, cols, self.z):
            self.doc_topic_counts[i] += zvec
            self.word_topic_counts[j] += zvec

    def generate(self, N, keep=True):
        word_probs = np.dot(self.theta, self.beta.T)
        assert word_probs.shape == (self.D, self.V)
        data = np.zeros((self.D, self.V))
        for d in range(self.D):
            data[d] = np.random.multinomial(N, word_probs[d])
        data = scipy.sparse.csr_matrix(data)

        if keep:
            self.data = data
            self.z = np.zeros((data.data.shape[0], self.T), dtype='uint32')
            self.resample_z()
            # precompute
            self._training_gammalns = \
                gammaln(data.sum(1)+1).sum() - gammaln(data.data+1).sum()

class StandardLDA(_LDABase):
    "Standard LDA with Dirichlet priors"

    def __init__(self, data, T, alpha_beta, alpha_theta):
        self.alpha_theta = alpha_theta
        super(StandardLDA, self).__init__(data, T, alpha_beta)

    @property
    def theta(self):
        return self._theta

    @theta.setter
    def theta(self, theta):
        self._theta = theta

    def initialize_theta(self):
        self.theta = sample_dirichlet(
            self.alpha_theta * np.ones((self.D, self.T)), 'horiz')

    def resample_theta(self):
        self.theta = sample_dirichlet(
            self.alpha_theta + self.doc_topic_counts, 'horiz')

    def copy_sample(self):
        new = copy.copy(self)
        new.beta = self.beta.copy()
        new._theta = self._theta.copy()
        return new

    def resample_collapsed(self,niter=1):
        self.resample_z_collapsed(niter)
        self.resample_theta()
        self.resample_beta()

    def resample_z_collapsed(self,niter=1):
        from ._lda import CollapsedCounts

        counts = CollapsedCounts(
            self.alpha_theta, self.alpha_beta, self.T,
            self.z, self.doc_topic_counts, self.word_topic_counts,
            self.data)
        counts.resample(niter)

        self.z = counts.z
        self.doc_topic_counts = counts.doc_topic_counts
        self.word_topic_counts = counts.word_topic_counts

###
# Correlated LDA Models (CTMs)
###

class StickbreakingCorrelatedLDA(_LDABase):
    "Correlated LDA with the stick breaking representation"

    def __init__(self, data, T, alpha_beta):
        mu, sigma = compute_uniform_mean_psi(T)
        self.theta_prior = Gaussian(
            mu=mu, sigma=sigma, mu_0=mu, sigma_0=T*sigma/10.,
            nu_0=T/10., kappa_0=1./10)

        self.ppgs = initialize_polya_gamma_samplers()
        self.omega = np.zeros((data.shape[0], T-1))

        super(StickbreakingCorrelatedLDA, self).__init__(data, T, alpha_beta)

    @property
    def theta(self):
        return psi_to_pi(self.psi)

    @theta.setter
    def theta(self, theta):
        self.psi = pi_to_psi(theta)

    def initialize_theta(self):
        self.psi = np.tile(self.theta_prior.mu, (self.D, 1))

    def resample_theta(self):
        self.resample_omega()
        self.resample_psi()

    def resample(self):
        super(StickbreakingCorrelatedLDA, self).resample()
        self.resample_theta_prior()

    def resample_omega(self):
        pgdrawvpar(
            self.ppgs, N_vec(self.doc_topic_counts).astype('float64').ravel(),
            self.psi.ravel(), self.omega.ravel())
        np.clip(self.omega, 1e-32, np.inf, out=self.omega)

    def resample_psi(self):
        Lmbda = np.linalg.inv(self.theta_prior.sigma)
        h = Lmbda.dot(self.theta_prior.mu)
        randvec = np.random.randn(self.D, self.T-1)  # pre-generate randomness

        for d, c in enumerate(self.doc_topic_counts):
            self.psi[d] = sample_infogaussian(
                Lmbda + np.diag(self.omega[d]), h + kappa_vec(c),
                randvec[d])

    def resample_theta_prior(self):
        self.theta_prior.resample(self.psi)

    def copy_sample(self):
        new = copy.copy(self)
        new.beta = self.beta.copy()
        new.psi = self.psi.copy()
        new.theta_prior = self.theta_prior.copy_sample()
        del new.z
        del new.omega
        return new

class LogisticNormalCorrelatedLDA(_LDABase):
    "Correlated LDA with the stick breaking representation"

    def __init__(self, data, T, alpha_beta):
        mu, sigma = np.zeros(T), np.eye(T)
        self.theta_prior = \
            Gaussian(
                mu=mu, sigma=sigma, mu_0=mu, sigma_0=T*sigma/10.,
                nu_0=T/10., kappa_0=10.)

        self.ppgs = initialize_polya_gamma_samplers()
        self.omega = np.zeros((data.shape[0], T))

        super(LogisticNormalCorrelatedLDA, self).__init__(data, T, alpha_beta)

    @property
    def theta(self):
        return ln_psi_to_pi(self.psi)

    @theta.setter
    def theta(self, theta):
        self.psi = ln_pi_to_psi(theta)

    def initialize_theta(self):
        self.psi = np.tile(self.theta_prior.mu, (self.D, 1))

    def resample_theta(self):
        self.resample_psi_and_omega()

    def resample(self):
        super(LogisticNormalCorrelatedLDA, self).resample()
        self.resample_theta_prior()

    def resample_psi_and_omega(self):
        Lmbda = np.linalg.inv(self.theta_prior.sigma)
        for d in range(self.D):
            N = self.data[d].sum()
            c = self.doc_topic_counts[d]
            for t in range(self.T):
                self.omega[d,t] = self.ppgs[0].pgdraw(
                    N, self._conditional_omega(d,t))

                mu_cond, sigma_cond = self._conditional_psi(d, t, Lmbda, N, c)
                self.psi[d,t] = np.random.normal(mu_cond, np.sqrt(sigma_cond))

    def _conditional_psi(self, d, t, Lmbda, N, c):
        nott = np.arange(self.T) != t
        psi = self.psi[d]
        omega = self.omega[d]
        mu = self.theta_prior.mu

        zetat = logsumexp(psi[nott])

        mut_marg = mu[t] - 1./Lmbda[t,t] * Lmbda[t,nott].dot(psi[nott] - mu[nott])
        sigmat_marg = 1./Lmbda[t,t]

        sigmat_cond = 1./(omega[t] + 1./sigmat_marg)

        # kappa is the mean dot precision, i.e. the sufficient statistic of a Gaussian
        # therefore we can sum over datapoints
        kappa = (c[t] - N/2.0).sum()
        mut_cond = sigmat_cond * (kappa + mut_marg / sigmat_marg + omega[t]*zetat)

        return mut_cond, sigmat_cond

    def _conditional_omega(self, d, t):
        nott = np.arange(self.T) != t
        psi = self.psi[d]
        zetat = logsumexp(psi[nott])
        return psi[t] - zetat

    def resample_theta_prior(self):
        self.theta_prior.resample(self.psi)

    def copy_sample(self):
        new = copy.copy(self)
        new.beta = self.beta.copy()
        new.psi = self.psi.copy()
        new.theta_prior = self.theta_prior.copy_sample()
        del new.z
        del new.omega
        return new

###
# Dynamic LDA Models (DTMs)
###

class StickbreakingDynamicTopicsLDA(object):
    def __init__(self, data, timestamps, K, alpha_theta):
        assert isinstance(data, scipy.sparse.csr.csr_matrix)
        self.alpha_theta = alpha_theta
        self.D, self.V = data.shape
        self.K = K

        self.data = data

        self.timestamps = timestamps
        self.timeidx = self._get_timeidx(timestamps, data)
        self.T = self.timeidx.max() - self.timeidx.min() + 1

        self.ppgs = initialize_polya_gamma_samplers()
        self.pyrngs = initialize_pyrngs()

        self.initialize_parameters()

        self._training_gammalns = \
            gammaln(data.sum(1)+1).sum() - gammaln(data.data+1).sum()

    def initialize_parameters(self):
        self.sigmasq_states = 0.1  # TODO make this learned, init from hypers

        mean_psi = compute_uniform_mean_psi(self.V)[0][None,:,None]
        self.psi = np.tile(mean_psi, (self.T, 1, self.K))

        self.omega = np.zeros_like(self.psi)

        self.theta = sample_dirichlet(
            self.alpha_theta * np.ones((self.D, self.K)), 'horiz')

        self.z = np.zeros((self.data.data.shape[0], self.K), dtype='uint32')
        self.resample_z()

    def log_likelihood(self, data=None):
        if data is not None:
            return log_likelihood(
                data, self._get_wordprobs(
                    data, self._get_timeidx(self.timestamps, data)))
        else:
            # this version avoids recomputing the training gammalns
            wordprobs = self._get_wordprobs(self.data, self.timeidx)
            return np.sum(np.nan_to_num(np.log(wordprobs)) * self.data.data) \
                + self._training_gammalns

    @property
    def beta(self):
        return psi_to_pi(self.psi, axis=1)

    def resample(self):
        self.resample_z()
        self.resample_theta()
        self.resample_beta()
        self.resample_lds_params()

    def resample_theta(self):
        self.theta = sample_dirichlet(
            self.alpha_theta + self.doc_topic_counts, 'horiz')

    def resample_beta(self):
        self.resample_omega()
        self.resample_psi()

    def resample_psi(self):
        mu_init, sigma_init, sigma_states, sigma_obs, y = \
            self._get_lds_effective_params()
        _, psi_flat = filter_and_sample_randomwalk(
            mu_init, sigma_init, sigma_states, sigma_obs, y)
        self.psi = psi_flat.reshape(self.psi.shape)

    def resample_z(self):
        topicprobs = self._get_topicprobs()
        multinomial(self.pyrngs, self.data.data, topicprobs, self.z)
        self._update_counts()

    def resample_omega(self):
        pgdrawvpar(
            self.ppgs,
            N_vec(self.time_word_topic_counts, axis=1)
                .astype('float64').ravel(),
            self.psi.ravel(), self.omega.ravel())
        np.clip(self.omega, 1e-32, np.inf, out=self.omega)

    def resample_lds_params(self):
        pass  # TODO

    def _get_lds_effective_params(self):
        mu_uniform, sigma_uniform = compute_uniform_mean_psi(self.V)
        mu_init = np.tile(mu_uniform, self.K)
        sigma_init = np.tile(np.diag(sigma_uniform), self.K)

        sigma_states = np.repeat(self.sigmasq_states, (self.V - 1) * self.K)

        sigma_obs = 1./self.omega
        y = kappa_vec(self.time_word_topic_counts, axis=1) / self.omega

        return mu_init, sigma_init, sigma_states, \
            sigma_obs.reshape(y.shape[0], -1), y.reshape(y.shape[0], -1)

    def _update_counts(self):
        self.doc_topic_counts = np.zeros((self.D, self.K), dtype='uint32')
        self.time_word_topic_counts = np.zeros((self.T, self.V, self.K), dtype='uint32')
        rows, cols = csr_nonzero(self.data)
        for i, j, t, zvec in zip(rows, cols, self.timeidx, self.z):
            self.doc_topic_counts[i] += zvec
            self.time_word_topic_counts[t,j] += zvec

    def _get_topicprobs(self):
        rows, cols = csr_nonzero(self.data)
        return normalize_rows(self.theta[rows] * self.beta[self.timeidx,cols])

    def _get_wordprobs(self, data, timeidx):
        rows, cols = csr_nonzero(data)
        return np.einsum('tk,tk->t',self.theta[rows],self.beta[timeidx, cols])

    def _get_timeidx(self, timestamps, data):
        timeidx = np.repeat(
            timeindices_from_timestamps(timestamps), np.diff(data.indptr))
        assert is_sorted(timeidx)
        return timeidx

# TODO we probably want to operate around a uniform (or separately sampled) bias
# point for psi. or maybe LDS should just learn a mean offset.