#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Sep 21 16:52:10 2017

@author: margauxmouchene
"""

from landlab import RasterModelGrid
from landlab.components import TransportLengthHillslopeDiffuser
from landlab.components import FlowDirectorSteepest, FlowAccumulator
import numpy as np
from numpy.testing import assert_equal, assert_almost_equal
import pytest

def test_route_to_multiple_error_raised():
    mg = RasterModelGrid((10, 10))
    z = mg.add_zeros('node', 'topographic__elevation')
    z += mg.x_of_node + mg.y_of_node
    fa = FlowAccumulator(mg, flow_director='MFD')
    fa.run_one_step()

    with pytest.raises(NotImplementedError):
        TransportLengthHillslopeDiffuser(mg, erodibility=1.0, slope_crit=0.5)


def test_tl_hill_diff():
    """Test cases where S>Sc, S=Sc and S<Sc"""

    # Test cases where S>Sc, S=Sc and S<Sc
    # Set up a 3x16 grid with closed boundaries and initial elevations.
    mg = RasterModelGrid((3, 12), 1.)
    z = np.array([0., 0., 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0., 0.,
                  0., 5., 1.9, 1.9, 1.9, 1.9, 1.3, 1.3, 1.3, 1.3, 1., 0.,
                  0., 0., 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0., 0.])
    _ = mg.add_field('node', 'topographic__elevation', z)
    mg.set_closed_boundaries_at_grid_edges(True, True, True, True)

    # Parameter values for test
    k = 0.001
    Sc = 0.6

    # Instantiate flow director and tl hillslope diffuser
    fdir = FlowDirectorSteepest(mg)
    tl_diff = TransportLengthHillslopeDiffuser(mg,
                                               erodibility=k,
                                               slope_crit=Sc)

    # Run flow director
    fdir.run_one_step()

    # test slopes
    s_out = mg.at_node['topographic__steepest_slope']
    s_test = np.array([0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,
                       0. ,  0. , 0. , 0. ,  3.1,  0. ,  0. ,  0. ,  0.6,  0. ,
                       0. ,  0. ,  0.3,  0. , 0. , 0. ,  0. ,  0. ,  0. ,  0. ,
                       0. ,  0. ,  0. ,  0. ,  0. ,  0. , 0. ])
    assert_almost_equal(s_out, s_test, decimal=10)

    # Run tl hillslope diffusion component
    tl_diff.run_one_step(1.)

    # Test results
    # flux_out
    fo_test = np.array([0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,
                       0. ,  0. , 0. , 0. ,  0.025,  0. ,  0. ,  0. ,  0.0006,
                       0. ,  0. ,  0. ,  0.0003,  0. , 0. , 0. ,  0. ,  0. ,
                       0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. , 0. ])
    fo_out = mg.at_node['sediment__flux_out']
    assert_almost_equal(fo_out, fo_test, decimal=10)

    # updated elevation
    elev_test = np.array([ 0.    ,  0.    ,  0.    ,  0.    ,  0.    ,  0.    ,
                          0.    ,0.    ,  0.    ,  0.    ,  0.    ,  0.    ,
                          0.    ,  4.975 ,  1.9   ,  1.9   ,  1.9   ,  1.8994,
                          1.3   , 1.3   ,  1.3   ,  1.2997,  1.    ,  0.    ,
                          0.    ,  0.    ,  0.    ,  0.    ,  0.    ,  0.    ,
                          0.    ,  0.    ,  0.    ,  0.    ,  0.    ,  0.    ])
    elev_out = mg.at_node['topographic__elevation']
    assert_almost_equal(elev_out, elev_test, decimal=10)

    # Run another time step because deposition and transfer were null
    # the first time
    fdir.run_one_step()
    tl_diff.run_one_step(1.)

    # Test results
    # flux_out
    fo_test = np.array([  0.00000000e+00,   0.00000000e+00,   0.00000000e+00,
                          0.00000000e+00,   0.00000000e+00,   0.00000000e+00,
                          0.00000000e+00,   0.00000000e+00,   0.00000000e+00,
                          0.00000000e+00,   0.00000000e+00,   0.00000000e+00,
                          0.00000000e+00,   2.47500000e-02,   0.00000000e+00,
                          0.00000000e+00,   6.00000000e-07,   5.99400000e-04,
                          0.00000000e+00,   0.00000000e+00,   3.00000000e-07,
                          2.99700000e-04,   0.00000000e+00,   0.00000000e+00,
                          0.00000000e+00,   0.00000000e+00,   0.00000000e+00,
                          0.00000000e+00,   0.00000000e+00,   0.00000000e+00,
                          0.00000000e+00,   0.00000000e+00,   0.00000000e+00,
                          0.00000000e+00,   0.00000000e+00,   0.00000000e+00])
    fo_out = mg.at_node['sediment__flux_out']
    assert_almost_equal(fo_out, fo_test, decimal=10)

    # updated elevation
    elev_test = np.array([ 0.       ,  0.       ,  0.       ,  0.       ,
                           0.       ,  0.       ,  0.       ,  0.       ,
                           0.       ,  0.       ,  0.       ,  0.       ,
                           0.       ,  4.95025  ,  1.925    ,  1.9      ,
                           1.8999994,  1.8988006,  1.3006   ,  1.3      ,
                           1.2999997,  1.2994003,  1.0003   ,  0.       ,
                           0.       ,  0.       ,  0.       ,  0.       ,
                           0.       ,  0.       ,  0.       ,  0.       ,
                           0.       ,  0.       , 0.       ,  0.       ])
    elev_out = mg.at_node['topographic__elevation']
    assert_almost_equal(elev_out, elev_test, decimal=10)
