"""Test potentiality flowrouter.

This tester turns over the potentiality flowrouter to ensure basic
functionality is working.

Probably doesn't account for cell area yet.
"""
import numpy as np
from numpy.testing import assert_allclose

from landlab import RasterModelGrid
from landlab.components import PotentialityFlowRouter

NROWS = 10
NCOLS = 10
DX = 1000.0


def test_sheetflow():
    flux = np.array(
        [
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
            1.00000000e06,
            1.87867966e06,
            2.72182541e06,
            3.54415588e06,
            4.35303038e06,
            5.15223689e06,
            5.94408409e06,
            6.73016503e06,
            6.51166745e06,
            1.00000000e00,
            1.00000000e06,
            2.12132034e06,
            3.24264069e06,
            4.35965005e06,
            5.47056275e06,
            6.57568517e06,
            7.67583899e06,
            8.77188984e06,
            8.86460677e06,
            1.00000000e00,
            1.00000000e06,
            2.00000000e06,
            3.03553391e06,
            4.08578644e06,
            5.14392129e06,
            6.20695360e06,
            7.27341175e06,
            8.34244102e06,
            8.41347211e06,
            1.00000000e00,
            1.00000000e06,
            2.00000000e06,
            3.00000000e06,
            4.01040764e06,
            5.03248558e06,
            6.06512434e06,
            7.10666517e06,
            8.15550411e06,
            8.21025766e06,
            1.00000000e00,
            1.00000000e06,
            2.00000000e06,
            3.00000000e06,
            4.01040764e06,
            5.03248558e06,
            6.06512434e06,
            7.10666517e06,
            8.15550411e06,
            8.21025766e06,
            1.00000000e00,
            1.00000000e06,
            2.00000000e06,
            3.03553391e06,
            4.08578644e06,
            5.14392129e06,
            6.20695360e06,
            7.27341175e06,
            8.34244102e06,
            8.41347211e06,
            1.00000000e00,
            1.00000000e06,
            2.12132034e06,
            3.24264069e06,
            4.35965005e06,
            5.47056275e06,
            6.57568517e06,
            7.67583899e06,
            8.77188984e06,
            8.86460677e06,
            1.00000000e00,
            1.00000000e06,
            1.87867966e06,
            2.72182541e06,
            3.54415588e06,
            4.35303038e06,
            5.15223689e06,
            5.94408409e06,
            6.73016503e06,
            6.51166745e06,
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
            1.00000000e00,
        ]
    )

    mg = RasterModelGrid((NROWS, NCOLS), xy_spacing=(DX, DX))
    z = (3000.0 - mg.node_x) * 0.5
    mg.at_node["topographic__elevation"] = z

    mg.set_closed_boundaries_at_grid_edges(False, True, True, True)
    mg.add_ones("water__unit_flux_in", at="node")

    pfr = PotentialityFlowRouter(mg)
    pfr.run_one_step()

    assert_allclose(mg.at_node["surface_water__discharge"], flux)


def test_in_network():
    # a valley network produced by stream power:
    z = np.array(
        [
            3.12900830e-04,
            3.66462671e-04,
            9.44438515e-04,
            1.45006772e-04,
            1.91406099e-04,
            4.42053204e-04,
            2.99818052e-04,
            5.45267467e-04,
            4.12129514e-04,
            7.43816953e-04,
            1.59251681e-04,
            7.39577249e01,
            6.68718419e01,
            2.95535987e01,
            7.69715256e01,
            4.61083179e01,
            5.70611522e01,
            6.99664564e01,
            8.29355817e01,
            5.85563532e-04,
            2.60300016e-04,
            8.99636726e01,
            1.11946320e02,
            5.16553709e01,
            1.38432251e02,
            1.02903579e02,
            1.27601424e02,
            8.59469601e01,
            4.77495429e01,
            8.00781223e-04,
            4.13095981e-04,
            7.41309307e01,
            1.31747968e02,
            1.11724617e02,
            1.71687943e02,
            1.93482716e02,
            1.61044205e02,
            1.34775824e02,
            7.63077925e01,
            8.85837646e-04,
            6.94907676e-04,
            5.53654246e01,
            9.95405009e01,
            1.81345288e02,
            1.88975196e02,
            1.78104180e02,
            1.75231160e02,
            1.14932425e02,
            5.12018382e01,
            1.26501797e-04,
            8.34527300e-04,
            8.80562940e01,
            1.24009142e02,
            1.55635807e02,
            1.56683637e02,
            1.62398410e02,
            1.04070282e02,
            6.99665030e01,
            6.41256897e01,
            5.50951003e-04,
            2.02545919e-04,
            3.05799760e01,
            5.02285119e01,
            8.31442034e01,
            1.75898080e02,
            1.80095770e02,
            1.23416767e02,
            6.97604901e01,
            3.04864568e01,
            4.42047775e-04,
            5.75778629e-04,
            8.13990441e01,
            1.17362500e02,
            1.35317452e02,
            1.78842796e02,
            9.65391990e01,
            1.16520146e02,
            1.59289373e02,
            8.29784784e01,
            4.21911459e-04,
            7.83145812e-04,
            8.44510235e01,
            6.85512072e01,
            3.99258990e01,
            9.07740020e01,
            4.47116751e01,
            1.05333999e02,
            1.04965376e02,
            6.42615041e01,
            8.48184002e-04,
            9.92952214e-04,
            5.44805365e-04,
            6.83657298e-04,
            4.27967811e-04,
            4.40101095e-04,
            5.47248461e-04,
            5.77178429e-06,
            4.39642103e-04,
            4.80194778e-04,
            9.24014550e-04,
        ]
    )

    flux = np.array(
        [
            0.71259633,
            1.8368437,
            7.98879652,
            9.29556375,
            10.55574455,
            7.24717392,
            7.35872882,
            4.50198662,
            1.97243346,
            0.54930274,
            1.71629518,
            4.0949232,
            4.50641349,
            19.45300944,
            5.73451891,
            10.2200909,
            8.84475793,
            5.62433494,
            3.63816461,
            3.2288009,
            3.64720491,
            4.08050742,
            3.55208345,
            12.80703452,
            4.51455295,
            7.22640241,
            4.82518478,
            6.0364361,
            8.21399889,
            4.89790743,
            6.37085345,
            7.20139071,
            3.66063349,
            6.65749537,
            4.5201495,
            3.16880878,
            4.56500213,
            4.47976874,
            5.91287906,
            7.65823897,
            6.75481467,
            10.33585975,
            7.95868492,
            3.44967745,
            3.30272455,
            3.79822256,
            3.72583825,
            5.35268783,
            11.33346169,
            7.61218948,
            13.23332132,
            5.52663873,
            4.71401637,
            5.25380598,
            6.32221242,
            5.21692832,
            9.95021529,
            12.12012981,
            9.62887235,
            12.51966215,
            12.20119879,
            28.07778172,
            21.69415422,
            13.44016311,
            3.42523128,
            3.16880878,
            5.79132165,
            11.8326263,
            20.62850507,
            10.58826741,
            10.25347039,
            3.76593706,
            3.58815683,
            4.10381596,
            3.23987285,
            7.32040184,
            4.62625103,
            3.16880878,
            4.05801502,
            8.29763076,
            1.42053684,
            3.54385495,
            4.9815377,
            7.32862308,
            5.884002,
            12.71539344,
            4.14640372,
            4.17838374,
            4.74896355,
            1.84587995,
            0.57901706,
            1.94727505,
            4.37381493,
            5.18582678,
            7.70532408,
            6.86738548,
            6.04728603,
            2.94714791,
            2.00238741,
            0.88296836,
        ]
    )

    potnt = np.array(
        [
            7.12596329e23,
            1.83684370e24,
            7.98879652e24,
            9.29556375e24,
            1.05557446e25,
            7.24717392e24,
            7.35872882e24,
            4.50198662e24,
            1.97243346e24,
            5.49302741e23,
            1.71629518e24,
            3.11607783e00,
            4.55006629e00,
            4.21948454e01,
            4.51089464e00,
            1.77476721e01,
            1.18679201e01,
            5.93094410e00,
            2.26828834e00,
            3.22880090e24,
            3.64720491e24,
            3.44492074e00,
            2.99141371e00,
            8.61459830e01,
            3.22150654e00,
            1.30493764e01,
            4.51470413e00,
            1.01713137e01,
            1.40167801e01,
            4.89790743e24,
            6.37085345e24,
            8.30460195e00,
            2.72370114e00,
            1.97019871e01,
            4.81322624e00,
            2.37203116e00,
            4.76156755e00,
            3.84190281e00,
            5.53509619e00,
            7.65823897e24,
            6.75481467e24,
            1.63796499e01,
            1.83238435e01,
            2.62626616e00,
            3.21546392e00,
            5.36483147e00,
            3.18817285e00,
            5.90189403e00,
            1.86765515e01,
            7.61218948e24,
            1.32333213e25,
            4.00510836e00,
            3.72650975e00,
            5.71530086e00,
            2.42783091e01,
            1.07981219e01,
            2.92283301e01,
            3.24892464e01,
            9.86359779e00,
            1.25196622e25,
            1.22011988e25,
            5.98718495e01,
            1.54766737e02,
            7.40803538e01,
            3.21059221e00,
            2.38746227e00,
            7.36767180e00,
            4.52838697e01,
            4.40547155e01,
            1.05882674e25,
            1.02534704e25,
            3.05098779e00,
            2.45849055e00,
            3.12955438e00,
            1.87486164e00,
            2.51124463e01,
            6.21239889e00,
            1.64106832e00,
            3.28554137e00,
            8.29763076e24,
            1.42053684e24,
            2.36944471e00,
            5.71708702e00,
            1.36764251e01,
            4.71435055e00,
            2.24231064e01,
            3.41364652e00,
            3.49569086e00,
            4.14217801e00,
            1.84587995e24,
            5.79017064e23,
            1.94727505e24,
            4.37381493e24,
            5.18582678e24,
            7.70532408e24,
            6.86738548e24,
            6.04728603e24,
            2.94714791e24,
            2.00238741e24,
            8.82968356e23,
        ]
    )

    mg = RasterModelGrid((NROWS, NCOLS), xy_spacing=(DX, DX))

    mg.add_field("topographic__elevation", z, at="node")

    Qin = np.ones_like(z) * 100.0 / (60.0 * 60.0 * 24.0 * 365.25)
    # ^remember, flux is /s, so this is a small number!
    mg.add_field("water__unit_flux_in", Qin, at="node")

    pfr = PotentialityFlowRouter(mg, flow_equation="Manning")
    pfr.run_one_step()

    assert_allclose(mg.at_node["surface_water__discharge"], flux)
    assert_allclose(mg.at_node["flow__potential"][mg.core_nodes], potnt[mg.core_nodes])
