import copy as cp

import numpy as np
import pytest
from numpy import testing

from landlab import FieldError, HexModelGrid, RasterModelGrid
from landlab.components import (
    FlowAccumulator,
    PriorityFloodFlowRouter,
    SpaceLargeScaleEroder,
)

def test_inputFields_flowRouter():
    """
    SpaceLargeScaleEroder should throw an error when topograhy is not equal to the sum of
    bedrock and soil thickness
    """
    # %%
    # Make a raster model grid and create a plateau
    mg = RasterModelGrid((5, 5))
    z = mg.add_zeros("topographic__elevation", at="node")
    br = mg.add_zeros("bedrock__elevation", at="node")
    br[:] = mg.x_of_node + mg.y_of_node
    soil = mg.add_zeros("soil__depth", at="node")
    z[:] = br + soil
    fa = FlowAccumulator(mg, flow_director="D8")
    fa.run_one_step()

    # make plateau at 10m
    br += 10

    # Instanciate the slider
    with pytest.raises(AssertionError):
        _ = SpaceLargeScaleEroder(mg)

# %%
def test_inputFields_soil():
    """
    SpaceLargeScaleEroder should throw an error when the soil__depth field is not provided
    """
    # %%
    mg = RasterModelGrid((5, 5))
    _ = mg.add_zeros("topographic__elevation", at="node")
    _ = mg.add_zeros("bedrock__elevation", at="node")
    fa = FlowAccumulator(mg, flow_director="D8")
    fa.run_one_step()

    # Instanciate the slider
    with pytest.raises(FieldError):
        _ = SpaceLargeScaleEroder(mg)

# %%
def test_inputFields_bedrock():
    """
    SpaceLargeScaleEroder should instanciate the bedrock__elevation field
    when it is not provided
    """
    # %%
    mg = RasterModelGrid((5, 5))
    _ = mg.add_zeros("topographic__elevation", at="node")
    _ = mg.add_zeros("soil__depth", at="node")
    fa = FlowAccumulator(mg, flow_director="D8")
    fa.run_one_step()
    _ = SpaceLargeScaleEroder(mg)

    assert "bedrock__elevation" in mg.at_node.keys()

# %%
def test_properties_phi_fraction_fines_LS():
    """
    SpaceLargeScaleEroder should throw an error when phi/fraction_fines_LS < 0 or phi > 0
    """
    # %%
    mg = RasterModelGrid((5, 5))
    z = mg.add_zeros("topographic__elevation", at="node")
    br = mg.add_zeros("bedrock__elevation", at="node")
    br[:] = mg.x_of_node + mg.y_of_node
    soil = mg.add_zeros("soil__depth", at="node")
    z[:] = br + soil
    fa = FlowAccumulator(mg, flow_director="D8")
    fa.run_one_step()

    # Instanciate the slider
    with pytest.raises(ValueError):
        _ = SpaceLargeScaleEroder(mg, phi=-0.2)
    # Instanciate the slider
    with pytest.raises(ValueError):
        _ = SpaceLargeScaleEroder(mg, phi=1.2)
    # Instanciate the slider
    with pytest.raises(ValueError):
        _ = SpaceLargeScaleEroder(mg, F_f=-0.2)
    # Instanciate the slider
    with pytest.raises(ValueError):
        _ = SpaceLargeScaleEroder(mg, F_f=1.2)

# %%

def test_route_to_multiple_error_raised():
    # %%
    mg = RasterModelGrid((10, 10))
    z = mg.add_zeros("topographic__elevation", at="node")
    br = mg.add_zeros("bedrock__elevation", at="node")
    br[:] = mg.x_of_node + mg.y_of_node
    soil = mg.add_zeros("soil__depth", at="node")
    z[:] = br + soil
    fa = FlowAccumulator(mg, flow_director="MFD")
    fa.run_one_step()

    with pytest.raises(NotImplementedError):
        SpaceLargeScaleEroder(
            mg,
            K_sed=0.1,
            K_br=0.1,
            F_f=0.5,
            phi=0.1,
            H_star=1.0,
            v_s=0.001,
            m_sp=1.0,
            n_sp=0.5,
            sp_crit_sed=0,
            sp_crit_br=0,
        )

# %%
def test_soil_field_already_on_grid():
    # %%
    """
    Test that an existing soil grid field is not changed by instantiating
    SpaceLargeScaleEroder.
    """

    # set up a 5x5 grid with one open outlet node and low initial elevations.
    nr = 5
    nc = 5
    mg = RasterModelGrid((nr, nc), xy_spacing=10.0)

    z = mg.add_zeros("topographic__elevation", at="node")
    br = mg.add_zeros("bedrock__elevation", at="node")
    soil = mg.add_zeros("soil__depth", at="node")
    soil += 1.0  # add 1m of soil everywehre

    mg["node"]["topographic__elevation"] += (
        mg.node_y / 10000 + mg.node_x / 10000 + np.random.rand(len(mg.node_y)) / 10000
    )
    mg.set_closed_boundaries_at_grid_edges(
        bottom_is_closed=True,
        left_is_closed=True,
        right_is_closed=True,
        top_is_closed=True,
    )
    mg.set_watershed_boundary_condition_outlet_id(
        0, mg["node"]["topographic__elevation"], -9999.0
    )
    br[:] = z - soil

    # Create a D8 flow handler
    FlowAccumulator(mg, flow_director="D8")

    # Instantiate SpaceLargeScaleEroder
    sp = SpaceLargeScaleEroder(
        mg,
        K_sed=0.01,
        K_br=0.01,
        F_f=0.0,
        phi=0.0,
        v_s=0.001,
        m_sp=0.5,
        n_sp=1.0,
        sp_crit_sed=0,
        sp_crit_br=0,
    )

    # ensure that 'soil__depth' field is everywhere equal to 1.0 m.
    testing.assert_array_equal(
        np.ones(mg.number_of_nodes),
        sp._soil__depth,
        err_msg="SpaceLargeScaleEroder soil depth field test failed",
        verbose=True,
    )

    # %% Check getters
    testing.assert_array_equal(
        0.01,
        sp.K_br,
        err_msg="Parameter value issue",
        verbose=True,
    )
    testing.assert_array_equal(
        0.01,
        sp.K_sed,
        err_msg="Parameter value issue",
        verbose=True,
    )
    # sediment erosion is zero before running the component
    testing.assert_array_equal(
        np.zeros(mg.number_of_nodes),
        sp.Es,
        err_msg="Parameter value issue",
        verbose=True,
    )
    # rock erosion is zero before running the component
    testing.assert_array_equal(
        np.zeros(mg.number_of_nodes),
        sp.Er,
        err_msg="Parameter value issue",
        verbose=True,
    )
    # %% Check setters
    sp.K_br = 0.02
    testing.assert_array_equal(
        0.02,
        sp.K_br,
        err_msg="Parameter value issue",
        verbose=True,
    )
    sp.K_sed = 0.02
    testing.assert_array_equal(
        0.02,
        sp.K_sed,
        err_msg="Parameter value issue",
        verbose=True,
    )

    with pytest.raises(AttributeError):
        sp.Es = np.zeros(mg.number_of_nodes)

    with pytest.raises(AttributeError):
        sp.Er = np.zeros(mg.number_of_nodes)

# %%

def test_br_field_already_on_grid():
    # %%
    """
    Test that an existing bedrock elevation grid field is not changed by
    instantiating SpaceLargeScaleEroder.
    """

    # set up a 5x5 grid with one open outlet node and low initial elevations.
    nr = 5
    nc = 5
    mg = RasterModelGrid((nr, nc), xy_spacing=10.0)

    z = mg.add_zeros("topographic__elevation", at="node")
    br = mg.add_zeros("bedrock__elevation", at="node")
    br += 1.0  # make bedrock elevation 5m below surface
    soil = mg.add_zeros("soil__depth", at="node")

    mg["node"]["topographic__elevation"] += (
        mg.node_y / 10000 + mg.node_x / 10000 + np.random.rand(len(mg.node_y)) / 10000
    )
    mg.set_closed_boundaries_at_grid_edges(
        bottom_is_closed=True,
        left_is_closed=True,
        right_is_closed=True,
        top_is_closed=True,
    )
    mg.set_watershed_boundary_condition_outlet_id(
        0, mg["node"]["topographic__elevation"], -9999.0
    )
    z[:] = br + soil

    # Create a D8 flow handler
    FlowAccumulator(mg, flow_director="D8")

    # Instantiate SpaceLargeScaleEroder
    sp = SpaceLargeScaleEroder(
        mg,
        K_sed=0.01,
        K_br=0.01,
        F_f=0.0,
        phi=0.0,
        v_s=0.001,
        m_sp=0.5,
        n_sp=1.0,
        sp_crit_sed=0,
        sp_crit_br=0,
    )

    # ensure that 'bedrock__elevation' field is everywhere equal to 1.0 m.
    testing.assert_array_equal(
        np.ones(mg.number_of_nodes),
        sp._bedrock__elevation,
        err_msg="SpaceLargeScaleEroder bedrock field test failed",
        verbose=True,
    )

# %%
def test_matches_detachment_solution():
    # %%
    """
    Test that model matches the detachment-limited analytical solution
    for slope/area relationship at steady state: S=(U/K_br)^(1/n)*A^(-m/n).
    """

    # %% set up a 5x5 grid with one open outlet node and low initial elevations.
    nr = 5
    nc = 5
    mg = RasterModelGrid((nr, nc), xy_spacing=10.0)

    z = mg.add_zeros("topographic__elevation", at="node")
    br = mg.add_zeros("bedrock__elevation", at="node")
    soil = mg.add_zeros("soil__depth", at="node")

    mg["node"]["topographic__elevation"] += (
        mg.node_y / 10000 + mg.node_x / 10000 + np.random.rand(len(mg.node_y)) / 10000
    )
    mg.set_closed_boundaries_at_grid_edges(
        bottom_is_closed=True,
        left_is_closed=True,
        right_is_closed=True,
        top_is_closed=True,
    )
    mg.set_watershed_boundary_condition_outlet_id(
        0, mg["node"]["topographic__elevation"], -9999.0
    )
    br[:] = z[:] - soil[:]

    # Create a D8 flow handler
    fa = FlowAccumulator(mg, flow_director="D8")

    # Parameter values for detachment-limited test
    K_br = 0.01
    U = 0.0001
    dt = 1.0
    F_f = 1.0  # all detached rock disappears; detachment-ltd end-member
    m_sp = 0.5
    n_sp = 1.0

    # Instantiate the SpaceLargeScaleEroder component...
    sp = SpaceLargeScaleEroder(
        mg,
        K_sed=0.00001,
        K_br=K_br,
        F_f=F_f,
        phi=0.1,
        H_star=1.0,
        v_s=0.001,
        m_sp=m_sp,
        n_sp=n_sp,
        sp_crit_sed=0,
        sp_crit_br=0,
    )

    # ... and run it to steady state (2000x1-year timesteps).
    for _ in range(2000):
        fa.run_one_step()
        sp.run_one_step(dt=dt)
        z[mg.core_nodes] += U * dt  # m
        br[mg.core_nodes] = z[mg.core_nodes] - soil[mg.core_nodes]

    # compare numerical and analytical slope solutions
    num_slope = mg.at_node["topographic__steepest_slope"][mg.core_nodes]
    analytical_slope = np.power(U / K_br, 1.0 / n_sp) * np.power(
        mg.at_node["drainage_area"][mg.core_nodes], -m_sp / n_sp
    )

    # test for match with analytical slope-area relationship
    testing.assert_array_almost_equal(
        num_slope,
        analytical_slope,
        decimal=8,
        err_msg="SpaceLargeScaleEroder detachment-limited test failed",
        verbose=True,
    )

# %%
def test_matches_detachment_solution_n_gr_1():
    # %%
    """
    Test that model matches the detachment-limited analytical solution
    for slope/area relationship at steady state: S=(U/K_br)^(1/n)*A^(-m/n).
    """

    # %% set up a 5x5 grid with one open outlet node and low initial elevations.
    nr = 5
    nc = 5
    mg = RasterModelGrid((nr, nc), xy_spacing=10.0)

    z = mg.add_zeros("topographic__elevation", at="node")
    br = mg.add_zeros("bedrock__elevation", at="node")
    soil = mg.add_zeros("soil__depth", at="node")

    mg["node"]["topographic__elevation"] += (
        mg.node_y / 10000 + mg.node_x / 10000 + np.random.rand(len(mg.node_y)) / 10000
    )
    mg.set_closed_boundaries_at_grid_edges(
        bottom_is_closed=True,
        left_is_closed=True,
        right_is_closed=True,
        top_is_closed=True,
    )
    mg.set_watershed_boundary_condition_outlet_id(
        0, mg["node"]["topographic__elevation"], -9999.0
    )
    br[:] = z[:] - soil[:]

    # Create a D8 flow handler
    fa = FlowAccumulator(mg, flow_director="D8")

    # Parameter values for detachment-limited test
    K_br = 0.01
    U = 0.0001
    dt = 1.0
    F_f = 1.0  # all detached rock disappears; detachment-ltd end-member
    m_sp = 0.5
    n_sp = 1.1

    # Instantiate the SpaceLargeScaleEroder component...
    sp = SpaceLargeScaleEroder(
        mg,
        K_sed=0.00001,
        K_br=K_br,
        F_f=F_f,
        phi=0.1,
        H_star=1.0,
        v_s=0.001,
        m_sp=m_sp,
        n_sp=n_sp,
        sp_crit_sed=0,
        sp_crit_br=0,
    )

    # ... and run it to steady state (2000x1-year timesteps).
    for _ in range(4000):
        fa.run_one_step()
        sp.run_one_step(dt=dt)
        z[mg.core_nodes] += U * dt  # m
        br[mg.core_nodes] = z[mg.core_nodes] - soil[mg.core_nodes]

    # compare numerical and analytical slope solutions
    num_slope = mg.at_node["topographic__steepest_slope"][mg.core_nodes]
    analytical_slope = np.power(U / K_br, 1.0 / n_sp) * np.power(
        mg.at_node["drainage_area"][mg.core_nodes], -m_sp / n_sp
    )

    # test for match with analytical slope-area relationship
    testing.assert_array_almost_equal(
        num_slope,
        analytical_slope,
        decimal=8,
        err_msg="SpaceLargeScaleEroder detachment-limited test failed",
        verbose=True,
    )

# %%

@pytest.mark.slow
def test_matches_transport_solution():
    # %%
    """
    Test that model matches the transport-limited analytical solution
    for slope/area relationship at steady state: S=((U * v_s) / (K_sed * A^m)
    + U / (K_sed * A^m))^(1/n).

    Also test that model matches the analytical solution for steady-state
    sediment flux: Qs = U * A * (1 - phi).
    """

    # set up a 5x5 grid with one open outlet node and low initial elevations.
    nr = 5
    nc = 5
    mg = RasterModelGrid((nr, nc), xy_spacing=10.0)

    z = mg.add_zeros("topographic__elevation", at="node")
    br = mg.add_zeros("bedrock__elevation", at="node")
    soil = mg.add_zeros("soil__depth", at="node")

    mg["node"]["topographic__elevation"] += (
        mg.node_y / 100000 + mg.node_x / 100000 + np.random.rand(len(mg.node_y)) / 10000
    )
    mg.set_closed_boundaries_at_grid_edges(
        bottom_is_closed=True,
        left_is_closed=True,
        right_is_closed=True,
        top_is_closed=True,
    )
    mg.set_watershed_boundary_condition_outlet_id(
        0, mg["node"]["topographic__elevation"], -9999.0
    )
    soil[:] += 100.0  # initial soil depth of 100 m
    br[:] = z[:]
    z[:] += soil[:]

    # Create a D8 flow handler
    fa = FlowAccumulator(
        mg, flow_director="D8", depression_finder="DepressionFinderAndRouter"
    )

    # Parameter values for detachment-limited test
    K_sed = 0.01
    U = 0.0001
    dt = 1.0
    F_f = 1.0  # all detached rock disappears; detachment-ltd end-member
    m_sp = 0.5
    n_sp = 1.0
    v_s = 0.5
    phi = 0.5

    # Instantiate the SpaceLargeScaleEroder component...
    sp = SpaceLargeScaleEroder(
        mg,
        K_sed=K_sed,
        K_br=0.01,
        F_f=F_f,
        phi=phi,
        H_star=1.0,
        v_s=v_s,
        m_sp=m_sp,
        n_sp=n_sp,
        sp_crit_sed=0,
        sp_crit_br=0,
    )

    # ... and run it to steady state (5000x1-year timesteps).
    for _ in range(5000):
        fa.run_one_step()
        sp.run_one_step(dt=dt)
        br[mg.core_nodes] += U * dt  # m
        soil[
            0
        ] = 100.0  # enforce constant soil depth at boundary to keep lowering steady
        z[:] = br[:] + soil[:]

    # compare numerical and analytical slope solutions
    num_slope = mg.at_node["topographic__steepest_slope"][mg.core_nodes]
    analytical_slope = np.power(
        (
            (U * v_s * (1 - phi))
            / (K_sed * np.power(mg.at_node["drainage_area"][mg.core_nodes], m_sp))
        )
        + (
            (U * (1 - phi))
            / (K_sed * np.power(mg.at_node["drainage_area"][mg.core_nodes], m_sp))
        ),
        1.0 / n_sp,
    )

    # test for match with analytical slope-area relationship
    testing.assert_array_almost_equal(
        num_slope,
        analytical_slope,
        decimal=8,
        err_msg="SpaceLargeScaleEroder transport-limited slope-area test failed",
        verbose=True,
    )

    # compare numerical and analytical sediment flux solutions
    num_sedflux = mg.at_node["sediment__outflux"][mg.core_nodes]
    analytical_sedflux = U * mg.at_node["drainage_area"][mg.core_nodes] * (1 - phi)

    # test for match with anakytical sediment flux
    testing.assert_array_almost_equal(
        num_sedflux,
        analytical_sedflux,
        decimal=8,
        err_msg="SpaceLargeScaleEroder transport-limited sediment flux test failed",
        verbose=True,
    )
    # %%

@pytest.mark.slow
def test_matches_bedrock_alluvial_solution():
    # %%
    """
    Test that model matches the bedrock-alluvial analytical solution
    for slope/area relationship at steady state:
    S=((U * v_s * (1 - F_f)) / (K_sed * A^m) + U / (K_br * A^m))^(1/n).

    Also test that the soil depth everywhere matches the bedrock-alluvial
    analytical solution at steady state:
    H = -H_star * ln(1 - (v_s / (K_sed / (K_br * (1 - F_f)) + v_s))).
    """

    # set up a 5x5 grid with one open outlet node and low initial elevations.
    nr = 5
    nc = 5
    mg = RasterModelGrid((nr, nc), xy_spacing=10.0)

    z = mg.add_zeros("topographic__elevation", at="node")
    br = mg.add_zeros("bedrock__elevation", at="node")
    soil = mg.add_zeros("soil__depth", at="node")

    mg["node"]["topographic__elevation"] += (
        mg.node_y / 100000 + mg.node_x / 100000 + np.random.rand(len(mg.node_y)) / 10000
    )
    mg.set_closed_boundaries_at_grid_edges(
        bottom_is_closed=True,
        left_is_closed=True,
        right_is_closed=True,
        top_is_closed=True,
    )
    mg.set_watershed_boundary_condition_outlet_id(
        0, mg["node"]["topographic__elevation"], -9999.0
    )
    soil[:] += 0.0  # initial condition of no soil depth.
    br[:] = z[:]
    z[:] += soil[:]

    # Create a D8 flow handler
    fa = FlowAccumulator(
        mg, flow_director="D8", depression_finder="DepressionFinderAndRouter"
    )

    # Parameter values for detachment-limited test
    K_br = 0.002
    K_sed = 0.002
    U = 0.0001
    dt = 10.0
    F_f = 0.2  # all detached rock disappears; detachment-ltd end-member
    m_sp = 0.5
    n_sp = 1.0
    v_s = 0.25
    H_star = 0.1

    # Instantiate the SpaceLargeScaleEroder component...
    sp = SpaceLargeScaleEroder(
        mg,
        K_sed=K_sed,
        K_br=K_br,
        F_f=F_f,
        phi=0.0,
        H_star=H_star,
        v_s=v_s,
        m_sp=m_sp,
        n_sp=n_sp,
        sp_crit_sed=0,
        sp_crit_br=0,
    )

    # ... and run it to steady state (10000x1-year timesteps).
    for _ in range(10000):
        fa.run_one_step()
        sp.run_one_step(dt=dt)
        br[mg.core_nodes] += U * dt  # m
        soil[0] = 0.0  # enforce 0 soil depth at boundary to keep lowering steady
        z[:] = br[:] + soil[:]

    # compare numerical and analytical slope solutions
    num_slope = mg.at_node["topographic__steepest_slope"][mg.core_nodes]
    analytical_slope = np.power(
        (
            (U * v_s * (1 - F_f))
            / (K_sed * np.power(mg.at_node["drainage_area"][mg.core_nodes], m_sp))
        )
        + (U / (K_br * np.power(mg.at_node["drainage_area"][mg.core_nodes], m_sp))),
        1.0 / n_sp,
    )

    # test for match with analytical slope-area relationship
    testing.assert_array_almost_equal(
        num_slope,
        analytical_slope,
        decimal=8,
        err_msg="SpaceLargeScaleEroder bedrock-alluvial slope-area test failed",
        verbose=True,
    )

    # compare numerical and analytical soil depth solutions
    num_h = mg.at_node["soil__depth"][mg.core_nodes]
    analytical_h = -H_star * np.log(1 - (v_s / (K_sed / (K_br * (1 - F_f)) + v_s)))

    # test for match with analytical sediment depth
    testing.assert_array_almost_equal(
        num_h,
        analytical_h,
        decimal=5,
        err_msg="SpaceLargeScaleEroder bedrock-alluvial soil thickness test failed",
        verbose=True,
    )
    # %%

def test_can_run_with_hex():
    """Test that model can run with hex model grid."""
    # %%
    # Set up a 5x5 grid with open boundaries and low initial elevations.
    mg = HexModelGrid((7, 7))
    z = mg.add_zeros("topographic__elevation", at="node")
    _ = mg.add_zeros("soil__depth", at="node")
    z[:] = 0.01 * mg.x_of_node

    # Create a D8 flow handler
    fa = FlowAccumulator(mg, flow_director="FlowDirectorSteepest")

    # Parameter values for test 1
    U = 0.001
    dt = 10.0

    # Create the SpaceLargeScaleEroder component...
    sp = SpaceLargeScaleEroder(
        mg,
        K_sed=0.00001,
        K_br=0.00000000001,
        F_f=0.5,
        phi=0.1,
        H_star=1.0,
        v_s=0.001,
        m_sp=0.5,
        n_sp=1.0,
        sp_crit_sed=0,
        sp_crit_br=0,
    )

    # ... and run it to steady state.
    for _ in range(2000):
        fa.run_one_step()
        sp.run_one_step(dt=dt)
        z[mg.core_nodes] += U * dt

# %%
def test_matches_detachment_solution_PF():
    # %%
    """
    Test that model matches the detachment-limited analytical solution
    for slope/area relationship at steady state: S=(U/K_br)^(1/n)*A^(-m/n).
    """

    # set up a 5x5 grid with one open outlet node and low initial elevations.
    nr = 5
    nc = 5
    mg = RasterModelGrid((nr, nc), xy_spacing=10.0)

    z = mg.add_zeros("topographic__elevation", at="node")
    br = mg.add_zeros("bedrock__elevation", at="node")
    soil = mg.add_zeros("soil__depth", at="node")

    mg["node"]["topographic__elevation"] += (
        mg.node_y / 10000 + mg.node_x / 10000 + np.random.rand(len(mg.node_y)) / 10000
    )
    mg.set_closed_boundaries_at_grid_edges(
        bottom_is_closed=True,
        left_is_closed=True,
        right_is_closed=True,
        top_is_closed=True,
    )
    mg.set_watershed_boundary_condition_outlet_id(
        0, mg["node"]["topographic__elevation"], -9999.0
    )
    br[:] = z[:] - soil[:]

    fa = PriorityFloodFlowRouter(
        mg, surface="topographic__elevation", flow_metric="D8", suppress_out=True
    )
    fa.run_one_step()

    # Parameter values for detachment-limited test
    K_br = 0.01
    U = 0.0001
    dt = 1.0
    F_f = 1.0  # all detached rock disappears; detachment-ltd end-member
    m_sp = 0.5
    n_sp = 1.0

    # Instantiate the SpaceLargeScaleEroder component...
    sp = SpaceLargeScaleEroder(
        mg,
        K_sed=0.00001,
        K_br=K_br,
        F_f=F_f,
        phi=0.1,
        H_star=1.0,
        v_s=0.001,
        v_s_lake=0.001,
        m_sp=m_sp,
        n_sp=n_sp,
        sp_crit_sed=0,
        sp_crit_br=0,
    )

    # ... and run it to steady state (2000x1-year timesteps).
    for _ in range(2000):
        fa.run_one_step()
        sp.run_one_step(dt=dt)
        z[mg.core_nodes] += U * dt  # m
        br[mg.core_nodes] = z[mg.core_nodes] - soil[mg.core_nodes]

    # compare numerical and analytical slope solutions
    num_slope = mg.at_node["topographic__steepest_slope"][mg.core_nodes]
    analytical_slope = np.power(U / K_br, 1.0 / n_sp) * np.power(
        mg.at_node["drainage_area"][mg.core_nodes], -m_sp / n_sp
    )

    # test for match with analytical slope-area relationship
    testing.assert_array_almost_equal(
        num_slope,
        analytical_slope,
        decimal=8,
        err_msg="SpaceLargeScaleEroder detachment-limited test failed",
        verbose=True,
    )

# %%
@pytest.mark.slow
def test_matches_transport_solution_PF():

    """
    Test that model matches the transport-limited analytical solution
    for slope/area relationship at steady state: S=((U * v_s) / (K_sed * A^m)
    + U / (K_sed * A^m))^(1/n).

    Also test that model matches the analytical solution for steady-state
    sediment flux: Qs = U * A * (1 - phi).
    """
    # %%
    # set up a 5x5 grid with one open outlet node and low initial elevations.
    nr = 5
    nc = 5
    mg = RasterModelGrid((nr, nc), xy_spacing=10.0)

    z = mg.add_zeros("topographic__elevation", at="node")
    br = mg.add_zeros("bedrock__elevation", at="node")
    soil = mg.add_zeros("soil__depth", at="node")

    mg["node"]["topographic__elevation"] += (
        mg.node_y / 100000 + mg.node_x / 100000 + np.random.rand(len(mg.node_y)) / 10000
    )
    mg.set_closed_boundaries_at_grid_edges(
        bottom_is_closed=True,
        left_is_closed=True,
        right_is_closed=True,
        top_is_closed=True,
    )
    mg.set_watershed_boundary_condition_outlet_id(
        0, mg["node"]["topographic__elevation"], -9999.0
    )
    soil[:] += 100.0  # initial soil depth of 100 m
    br[:] = z[:]
    z[:] += soil[:]

    # Create a D8 flow handler
    fa = PriorityFloodFlowRouter(
        mg, surface="topographic__elevation", flow_metric="D8", suppress_out=True
    )
    fa.run_one_step()

    # Parameter values for detachment-limited test
    K_sed = 0.01
    U = 0.0001
    dt = 1.0
    F_f = 1.0  # all detached rock disappears; detachment-ltd end-member
    m_sp = 0.5
    n_sp = 1.0
    v_s = 0.5
    phi = 0.5

    # Instantiate the SpaceLargeScaleEroder component...
    sp = SpaceLargeScaleEroder(
        mg,
        K_sed=K_sed,
        K_br=0.01,
        F_f=F_f,
        phi=phi,
        H_star=1.0,
        v_s=v_s,
        m_sp=m_sp,
        n_sp=n_sp,
        sp_crit_sed=0,
        sp_crit_br=0,
    )

    # ... and run it to steady state (5000x1-year timesteps).
    for _ in range(5000):
        fa.run_one_step()
        sp.run_one_step(dt=dt)
        br[mg.core_nodes] += U * dt  # m
        soil[
            0
        ] = 100.0  # enforce constant soil depth at boundary to keep lowering steady
        z[:] = br[:] + soil[:]

    # compare numerical and analytical slope solutions
    num_slope = mg.at_node["topographic__steepest_slope"][mg.core_nodes]
    analytical_slope = np.power(
        (
            (U * v_s * (1 - phi))
            / (K_sed * np.power(mg.at_node["drainage_area"][mg.core_nodes], m_sp))
        )
        + (
            (U * (1 - phi))
            / (K_sed * np.power(mg.at_node["drainage_area"][mg.core_nodes], m_sp))
        ),
        1.0 / n_sp,
    )

    # test for match with analytical slope-area relationship
    testing.assert_array_almost_equal(
        num_slope,
        analytical_slope,
        decimal=8,
        err_msg="SpaceLargeScaleEroder transport-limited slope-area test failed",
        verbose=True,
    )

    # compare numerical and analytical sediment flux solutions
    num_sedflux = mg.at_node["sediment__outflux"][mg.core_nodes]
    analytical_sedflux = U * mg.at_node["drainage_area"][mg.core_nodes] * (1 - phi)

    # test for match with anakytical sediment flux
    testing.assert_array_almost_equal(
        num_sedflux,
        analytical_sedflux,
        decimal=8,
        err_msg="SpaceLargeScaleEroder transport-limited sediment flux test failed",
        verbose=True,
    )

# %%
@pytest.mark.slow
def test_matches_bedrock_alluvial_solution_PF():
    """
    Test that model matches the bedrock-alluvial analytical solution
    for slope/area relationship at steady state:
    S=((U * v_s * (1 - F_f)) / (K_sed * A^m) + U / (K_br * A^m))^(1/n).

    Also test that the soil depth everywhere matches the bedrock-alluvial
    analytical solution at steady state:
    H = -H_star * ln(1 - (v_s / (K_sed / (K_br * (1 - F_f)) + v_s))).
    """
    # %%
    # set up a 5x5 grid with one open outlet node and low initial elevations.
    nr = 5
    nc = 5
    mg = RasterModelGrid((nr, nc), xy_spacing=10.0)

    z = mg.add_zeros("topographic__elevation", at="node")
    br = mg.add_zeros("bedrock__elevation", at="node")
    soil = mg.add_zeros("soil__depth", at="node")

    mg["node"]["topographic__elevation"] += (
        mg.node_y / 100000 + mg.node_x / 100000 + np.random.rand(len(mg.node_y)) / 10000
    )
    mg.set_closed_boundaries_at_grid_edges(
        bottom_is_closed=True,
        left_is_closed=True,
        right_is_closed=True,
        top_is_closed=True,
    )
    mg.set_watershed_boundary_condition_outlet_id(
        0, mg["node"]["topographic__elevation"], -9999.0
    )
    soil[:] += 0.0  # initial condition of no soil depth.
    br[:] = z[:]
    z[:] += soil[:]

    # Create a D8 flow handler
    fa = PriorityFloodFlowRouter(
        mg, surface="topographic__elevation", flow_metric="D8", suppress_out=True
    )
    fa.run_one_step()

    # Parameter values for detachment-limited test
    K_br = 0.002
    K_sed = 0.002
    U = 0.0001
    dt = 10.0
    F_f = 0.2  # all detached rock disappears; detachment-ltd end-member
    m_sp = 0.5
    n_sp = 1.0
    v_s = 0.25
    H_star = 0.1

    # Instantiate the SpaceLargeScaleEroder component...
    sp = SpaceLargeScaleEroder(
        mg,
        K_sed=K_sed,
        K_br=K_br,
        F_f=F_f,
        phi=0.0,
        H_star=H_star,
        v_s=v_s,
        m_sp=m_sp,
        n_sp=n_sp,
        sp_crit_sed=0,
        sp_crit_br=0,
    )

    # ... and run it to steady state (10000x1-year timesteps).
    for _ in range(10000):
        fa.run_one_step()
        sp.run_one_step(dt=dt)
        br[mg.core_nodes] += U * dt  # m
        soil[0] = 0.0  # enforce 0 soil depth at boundary to keep lowering steady
        z[:] = br[:] + soil[:]

    # compare numerical and analytical slope solutions
    num_slope = mg.at_node["topographic__steepest_slope"][mg.core_nodes]
    analytical_slope = np.power(
        (
            (U * v_s * (1 - F_f))
            / (K_sed * np.power(mg.at_node["drainage_area"][mg.core_nodes], m_sp))
        )
        + (U / (K_br * np.power(mg.at_node["drainage_area"][mg.core_nodes], m_sp))),
        1.0 / n_sp,
    )

    # test for match with analytical slope-area relationship
    testing.assert_array_almost_equal(
        num_slope,
        analytical_slope,
        decimal=8,
        err_msg="SpaceLargeScaleEroder bedrock-alluvial slope-area test failed",
        verbose=True,
    )

    # compare numerical and analytical soil depth solutions
    num_h = mg.at_node["soil__depth"][mg.core_nodes]
    analytical_h = -H_star * np.log(1 - (v_s / (K_sed / (K_br * (1 - F_f)) + v_s)))

    # test for match with analytical sediment depth
    testing.assert_array_almost_equal(
        num_h,
        analytical_h,
        decimal=5,
        err_msg="SpaceLargeScaleEroder bedrock-alluvial soil thickness test failed",
        verbose=True,
    )
    # %%

def test_MassBalance():
    # %%
    # set up a 15x15 grid with one open outlet node and low initial elevations.
    nr = 15
    nc = 15
    mg = RasterModelGrid((nr, nc), xy_spacing=10.0)

    z = mg.add_zeros("topographic__elevation", at="node")
    br = mg.add_zeros("bedrock__elevation", at="node")
    soil = mg.add_zeros("soil__depth", at="node")

    mg["node"]["topographic__elevation"] += (
        mg.node_y / 100000 + mg.node_x / 100000 + np.random.rand(len(mg.node_y)) / 10000
    )
    mg.set_closed_boundaries_at_grid_edges(
        bottom_is_closed=True,
        left_is_closed=True,
        right_is_closed=True,
        top_is_closed=True,
    )
    mg.set_watershed_boundary_condition_outlet_id(
        0, mg["node"]["topographic__elevation"], -9999.0
    )
    soil[:] += 0.0  # initial condition of no soil depth.
    br[:] = z[:]
    z[:] += soil[:]

    # Create a D8 flow handler
    # fa = PriorityFloodFlowRouter(mg, surface="topographic__elevation", flow_metric = 'D8',suppress_out=True)
    # fa.run_one_step()

    # Create a D8 flow handler
    fa = FlowAccumulator(
        mg, flow_director="D8", depression_finder="DepressionFinderAndRouter"
    )

    # Parameter values for detachment-limited test
    K_br = 0.002
    K_sed = 0.002
    U = 0.0001
    dt = 10.0
    F_f = 0.2  # all detached rock disappears; detachment-ltd end-member
    m_sp = 0.5
    n_sp = 1.0
    v_s = 0.25
    H_star = 0.1

    # Instantiate the Space component...
    sp = SpaceLargeScaleEroder(
        mg,
        K_sed=K_sed,
        K_br=K_br,
        F_f=F_f,
        phi=0.0,
        H_star=H_star,
        v_s=v_s,
        m_sp=m_sp,
        n_sp=n_sp,
        sp_crit_sed=0,
        sp_crit_br=0,
    )
    # Get values before run
    z = mg.at_node["topographic__elevation"]
    br = mg.at_node["bedrock__elevation"]
    H = mg.at_node["soil__depth"]
    cores = mg.core_nodes
    area = mg.cell_area_at_node
    # ... and run it to steady state (10000x1-year timesteps).
    for _ in range(10000):
        fa.run_one_step()
        soil_B = cp.deepcopy(H)
        bed_B = cp.deepcopy(br)
        vol_SSY_riv, V_leaving_riv = sp.run_one_step(dt=dt)
        diff_MB = (
            np.sum((bed_B[cores] - br[cores]) * area[cores])
            + np.sum((soil_B[cores] - H[cores]) * area[cores]) * (1 - sp._phi)
            - vol_SSY_riv * dt
            - V_leaving_riv
        )

        br[mg.core_nodes] += U * dt  # m
        soil[0] = 0.0  # enforce 0 soil depth at boundary to keep lowering steady
        z[:] = br[:] + soil[:]

        # Test Every iteration
        testing.assert_array_almost_equal(
            z[cores],
            br[cores] + H[cores],
            decimal=5,
            err_msg="Topography does not equal sum of bedrock and soil! Decrease timestep",
            verbose=True,
        )
        testing.assert_array_less(
            abs(diff_MB),
            1e-8 * mg.number_of_nodes,
            err_msg="Mass balance error SpaceLargeScaleEroder! Try to resolve by becreasing timestep",
            verbose=True,
        )


from thesiaNet.tensor import tensor
import numpy as np

class Batch:
    def __init__(self,input, target, batch_size= 32, shuffle = True):
        """
        this class conver data into the batches which given size
        :param input: input of the model
        :param target: original output
        :param batch_size: bach size
        :param shuffle: Boolean values (True : shuffeling)
        """
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.starts = np.arange(0, len(input), self.batch_size)
        if shuffle:
            np.random.shuffle(self.starts)

        batch = []
        for start in self.starts:
            input_batch = input[start:start + self.batch_size]
            target_batch = target[start:start + self.batch_size]
            batch.append((input_batch, target_batch))

        self.batch = batch


"""
==============================================
The :mod:`mpi_array.globale_ufunc_test` Module
==============================================

Module defining :mod:`mpi_array.globale` unit-tests.
Execute as::

   python -m mpi_array.globale_ufunc_test

and with parallelism::

   mpirun -n  2 python -m mpi_array.globale_ufunc_test
   mpirun -n  4 python -m mpi_array.globale_ufunc_test
   mpirun -n 27 python -m mpi_array.globale_ufunc_test

Classes
=======

.. autosummary::
   :toctree: generated/
   :template: autosummary/inherits_TestCase_class.rst

   UfuncResultTypeTest - Tests for :func:`mpi_array.globale_ufunc.ufunc_result_type` function.
   BroadcastShapeTest - Tests for :func:`mpi_array.globale_ufunc.broadcast_shape` function.
   GndarrayUfuncTest - Tests for :func:`mpi_array.globale_ufunc.gndarray_array_ufunc` function.
   ToGndarrayConverter - Base class for :obj:`numpy.ndarray` to :obj:`mpi_array.globale.gndarray`.
"""
from __future__ import absolute_import

import numpy as _np

from .license import license as _license, copyright as _copyright, version as _version
from . import unittest as _unittest
from . import logging as _logging  # noqa: E402,F401
from .comms import LT_NODE, LT_PROCESS, DT_CLONED, DT_SINGLE_LOCALE, DT_BLOCK  # , DT_SLAB
from . import comms as _comms
from . import distribution as _distribution
from .globale_ufunc import broadcast_shape, ufunc_result_type, get_extents
from .globale_ufunc import check_equivalent_inter_locale_comms
from .globale import gndarray as _gndarray
from .globale_creation import ones as _ones, zeros as _zeros, asarray as _asarray
from .globale_creation import empty as _empty
from .globale import copyto as _copyto

__author__ = "Shane J. Latham"
__license__ = _license()
__copyright__ = _copyright()
__version__ = _version()

class UfuncResultTypeTest(_unittest.TestCase):

    """
    :obj:`unittest.TestCase` for :func:`mpi_array.globale_ufunc.ufunc_result_type`.
    """

    def test_single_output(self):
        """
        :obj:`unittest.TestCase` for :func:`mpi_array.globale_ufunc.ufunc_result_type`,
        with single output array.
        """

        rank_logger = _logging.get_rank_logger(self.id())

        uft = ['ff->f', 'll->l', 'cc->c', 'fl->b', 'dd->d']

        inputs = (_np.array([1, 2], dtype='l'), _np.array([3, 4], dtype='l'))
        outputs = None
        dtypes = ufunc_result_type(uft, inputs, outputs)
        rank_logger.debug("dtypes=%s", dtypes)
        self.assertSequenceEqual((_np.dtype('l'),), dtypes)

        inputs = (_np.array([1, 2], dtype='f'), _np.array([3, 4], dtype='f'))
        outputs = (_np.array([0, 0], dtype='f'),)
        dtypes = ufunc_result_type(uft, inputs, outputs)
        self.assertSequenceEqual((_np.dtype('f'),), dtypes)
        outputs = (_np.array([0, 0], dtype='d'),)
        dtypes = ufunc_result_type(uft, inputs, outputs)
        self.assertSequenceEqual((_np.dtype('d'),), dtypes)
        outputs = (_np.array([0, 0], dtype='b'),)
        self.assertRaises(ValueError, ufunc_result_type, uft, inputs, outputs)

        inputs = (_np.array([1, 2], dtype='f'), _np.array([3, 4], dtype='l'))
        outputs = None
        dtypes = ufunc_result_type(uft, inputs, outputs)
        self.assertSequenceEqual((_np.dtype('b'),), dtypes)

        inputs = (_np.array([1, 2], dtype='f'), 5.0)
        outputs = None
        dtypes = ufunc_result_type(uft, inputs, outputs)
        rank_logger.debug("dtypes=%s", dtypes)
        self.assertSequenceEqual((_np.dtype('f'),), dtypes)

        inputs = (_np.array([1, 2], dtype='f'), 5.0e150)
        outputs = None
        dtypes = ufunc_result_type(uft, inputs, outputs)
        rank_logger.debug("dtypes=%s", dtypes)
        self.assertSequenceEqual((_np.dtype('d'),), dtypes)

        inputs = (_np.array([1, 2], dtype='complex128'), 5.0e150)
        outputs = None
        self.assertRaises(
            ValueError,
            ufunc_result_type, uft, inputs, outputs
        )

    def test_tuple_input(self):
        """
        :obj:`unittest.TestCase` for :func:`mpi_array.globale_ufunc.ufunc_result_type`,
        with single output array.
        """

        rank_logger = _logging.get_rank_logger(self.id())

        uft = ['ff->f', 'll->l', 'cc->c', 'fl->b', 'dd->d']

        inputs = (_np.array((1, 2, 3, 4), dtype='int32'), (1, 2, 3, 4))
        outputs = None
        dtypes = ufunc_result_type(uft, inputs, outputs)
        rank_logger.debug("dtypes=%s", dtypes)
        self.assertSequenceEqual((_np.dtype('l'),), dtypes)

    def test_multiple_output(self):
        """
        :obj:`unittest.TestCase` for :func:`mpi_array.globale_ufunc.ufunc_result_type`,
        with multiple output arrays.
        """

        rank_logger = _logging.get_rank_logger(self.id())

        uft = ['eee->eBl', 'fff->fBl', 'ddd->dBl', ]

        inputs = (_np.array([1, 2], dtype='e'), _np.array([3, 4], dtype='f'), 4.0)
        outputs = None
        dtypes = ufunc_result_type(uft, inputs, outputs)
        rank_logger.debug("dtypes=%s", dtypes)
        self.assertSequenceEqual((_np.dtype('f'), _np.dtype('B'), _np.dtype('l')), dtypes)

        inputs = (_np.array([1, 2], dtype='e'), _np.array([3, 4], dtype='f'), 4.0)
        outputs = (_np.array([1, 2], dtype='f'),)
        dtypes = ufunc_result_type(uft, inputs, outputs)
        rank_logger.debug("dtypes=%s", dtypes)
        self.assertSequenceEqual((_np.dtype('f'), _np.dtype('B'), _np.dtype('l')), dtypes)

        inputs = (_np.array([1, 2], dtype='e'), _np.array([3, 4], dtype='f'), 4.0)
        outputs = (_np.array([1, 2], dtype='d'), _np.array([1, 2], dtype='i'))
        dtypes = ufunc_result_type(uft, inputs, outputs)
        rank_logger.debug("dtypes=%s", dtypes)
        self.assertSequenceEqual((_np.dtype('d'), _np.dtype('i'), _np.dtype('l')), dtypes)

        inputs = (_np.array([1, 2], dtype='e'), _np.array([3, 4], dtype='f'), 4.0)
        outputs = (_np.array([1, 2], dtype='d'), _np.array([1, 2], dtype='b'))
        self.assertRaises(
            ValueError,
            ufunc_result_type, uft, inputs, outputs
        )

        inputs = (_np.array([1, 2], dtype='e'), _np.array([3, 4], dtype='f'), 4.0)
        outputs = \
            (
                _np.array([1, 2], dtype='d'),
                _np.array([1, 2], dtype='i'),
                _np.array([1, 2], dtype='uint16')
            )
        self.assertRaises(
            ValueError,
            ufunc_result_type, uft, inputs, outputs
        )

    def test_example(self):
        import numpy as np
        import mpi_array as mpia
        try:
            inp = (
                np.zeros((10, 10, 10), dtype='float16'),
                16.0,
                mpia.zeros((10, 10, 10), dtype='float32'),
            )
            dtypes = ufunc_result_type(['eee->e?', 'fff->f?', 'ddd->d?'], inputs=inp)
            self.assertSequenceEqual((_np.dtype('float32'), _np.dtype('bool')), dtypes)
            out = (mpia.zeros((10, 10, 10), dtype="float64"),)
            dtypes = ufunc_result_type(['eee->e?', 'fff->f?', 'ddd->d?'], inputs=inp, outputs=out)
            self.assertSequenceEqual((_np.dtype('float64'), _np.dtype('bool')), dtypes)
            out += (mpia.zeros((10, 10, 10), dtype="uint16"),)
            dtypes = ufunc_result_type(['eee->e?', 'fff->f?', 'ddd->d?'], inputs=inp, outputs=out)
            self.assertSequenceEqual((_np.dtype('float64'), _np.dtype('uint16')), dtypes)
        finally:
            inp[2].free()
            out[0].free()
            out[1].free()

class BroadcastShapeTest(_unittest.TestCase):

    """
    :obj:`unittest.TestCase` for :func:`mpi_array.globale_ufunc.broadcast_shape`.
    """

    def test_non_broadcastable(self):
        """
        Test that :func:`mpi_array.globale_ufunc.broadcast_shape` raises
        a :obj:`ValueError` if shapes are not broadcastable.
        """

        self.assertRaises(
            ValueError,
            broadcast_shape,
            (4,),
            (5,)
        )
        self.assertRaises(
            ValueError,
            broadcast_shape,
            (5,),
            (6,)
        )
        self.assertRaises(
            ValueError,
            broadcast_shape,
            (5, 5),
            (6, 5)
        )
        self.assertRaises(
            ValueError,
            broadcast_shape,
            (5, 6),
            (5, 5)
        )
        self.assertRaises(
            ValueError,
            broadcast_shape,
            (5, 6, 7),
            (5, 1, 1),
            (5, 6, 1),
            (1, 1, 7),
            (1, 1, 6)
        )

    def test_broadcastable(self):
        """
        Asserts for variety of broadcastable shapes.
        """
        self.assertSequenceEqual((), broadcast_shape(()))
        self.assertSequenceEqual((), broadcast_shape((), ()))
        self.assertSequenceEqual((), broadcast_shape((), (), ()))

        self.assertSequenceEqual((1,), broadcast_shape((), (), (1, )))
        self.assertSequenceEqual((4, ), broadcast_shape((4, ), (), (1, )))
        self.assertSequenceEqual((1, 4), broadcast_shape((4, ), (1, 1), (1, )))
        self.assertSequenceEqual((4, 5), broadcast_shape((5, ), (1, 5), (4, 1)))

class ToGndarrayConverter(object):

    """
    Base class for converting :obj:`numpy.ndarray` objects
    to :obj:`mpi_array.globale.gndarray` objects.
    """

    def __init__(self, **kwargs):
        """
        The :samp:`kwargs` are passed directly to the :func:`mpi_array.globale_creation.asarray`
        function in :meth:`__call__`.
        """
        self.kwargs = kwargs

    def __call__(self, npy_ary):
        """
        Converts the :samp:`{npy_ary}` to a :obj:`mpi_array.globale.gndarray` instance.

        :type npy_ary: :obj:`numpy.ndarray`
        :param npy_ary: Array converted to :obj:`mpi_array.globale.gndarray`.
           This array is assumed to be identical on all peer-rank MPI processes.
        :rtype: :obj:`mpi_array.globale.gndarray`
        :return: The :samp:`{npy_ary}` converted to a :obj:`mpi_array.globale.gndarray` instance.
        """
        if "halo" not in self.kwargs.keys():
            halo = _np.random.randint(low=1, high=4, size=(npy_ary.ndim, 2))
            gnd_ary = _asarray(npy_ary, halo=halo, **self.kwargs)
        else:
            gnd_ary = _asarray(npy_ary, **self.kwargs)

        return gnd_ary

class GetExtentsTest(_unittest.TestCase):

    """
    :obj:`unittest.TestCase` for :func:`mpi_array.globale_ufunc.get_extents`.
    """

    def setUp(self):
        """
        """
        locale_comms, inter_locale_rank_to_peer_rank, this_locale = _comms.create_locale_comms()
        self.comms = locale_comms
        self.locale_info = this_locale
        del inter_locale_rank_to_peer_rank

    def test_scalar(self):
        """
        Test :func:`mpi_array.globale_ufunc.get_extents` with scalar input.
        """
        l, g = get_extents(5.0, self.locale_info)
        self.assertTrue(isinstance(l, _distribution.ScalarLocaleExtent))
        self.assertTrue(isinstance(g, _distribution.ScalarGlobaleExtent))

class GndarrayUfuncTest(_unittest.TestCase):

    """
    :obj:`unittest.TestCase` for :obj:`mpi_array.globale_ufunc`.
    """

    def setUp(self):
        """
        Initialise :func:`numpy.random.seed`.
        """
        _np.random.seed(1531796312)
        self._rank_logger = _logging.get_rank_logger(self.id())
        with _asarray(_np.zeros((100,))) as gary:
            self.num_node_locales = gary.num_locales

    @property
    def rank_logger(self):
        """
        A :obj:`logging.Logger` object.
        """
        return self._rank_logger

    def compare_results(self, mpi_cln_npy_result_ary, mpi_result_ary):
        """
        Asserts that all elements of
        the :obj:`mpi_array.globale.gndarray` :samp:`{mpi_cln_npy_result_ary}`
        equal all elements of the  :obj:`mpi_array.globale.gndarray` :samp:`{mpi_result_ary}`.

        :type mpi_cln_npy_result_ary: :obj:`mpi_array.globale.gndarray`
        :param mpi_cln_npy_result_ary: The result returned by :samp:`{func}(*{func_args})`
            converted to a cloned-distribution :obj:`mpi_array.globale.gndarray`.
        :type mpi_result_ary: :obj:`mpi_array.globale.gndarray`
        :param mpi_result_ary: The result array
           from :meth:`mpi_array.globale.gndarray.__array_ufunc__` execution.
        """
        with \
                _zeros(
                    shape=mpi_result_ary.shape,
                    dtype=mpi_result_ary.dtype,
                    locale_type=LT_NODE,
                    distrib_type=DT_CLONED
                ) as mpi_cln_mpi_result_ary:

            _copyto(dst=mpi_cln_mpi_result_ary, src=mpi_result_ary)

            self.assertSequenceEqual(
                tuple(mpi_cln_npy_result_ary.shape),
                tuple(mpi_cln_mpi_result_ary.shape)
            )
            self.assertTrue(
                _np.all(
                    mpi_cln_npy_result_ary.lndarray_proxy.view_n
                    ==
                    mpi_cln_mpi_result_ary.lndarray_proxy.view_n
                )
            )

    def convert_func_args_to_gndarrays(self, converter, func_args):
        """

        :type converter: :obj:`ToGndarrayConverter`
        :param converter: Used to convert the :obj:`numpy.ndarray` instances
           of :samp:`{func_args}` to :obj:`mpi_array.globale.gndarray` instances.
        :type func_args: sequence of :obj:`numpy.ndarray` or array-like objects
        :param func_args: Sequence of array-like objects.
           Can be comprised of misture of :obj:`numpy.ndarray` instances, scalars
           or (broadcastable) sequences (e.g. tuple of scalars) elements.
        :rtype: :obj:`list`
        :return: The :samp:`{func_args}` list with :obj:`numpy.ndarray` instances
           converted to :obj:`mpi_array.globale.gndarray` instances.
        """
        return [converter(arg) if isinstance(arg, _np.ndarray) else arg for arg in func_args]

    def do_convert_execute_and_compare(self, mpi_cln_npy_result_ary, converter, func, *func_args):
        """
        Compares the result of :samp:`{func}` called
        with :samp:`self.convert_func_args_to_gndarrays({converter}, {func_args})`
        converted arguments with the :samp:`{mpi_cln_npy_result_ary}` array (which should
        have been produced by
        calling :samp:`mpi_array.globale_creation.asarray({func}(*func_args))`).

        :type mpi_cln_npy_result_ary: :obj:`mpi_array.globale.gndarray`
        :param mpi_cln_npy_result_ary: The result returned by :samp:`{func}(*{func_args})`
            converted to a cloned-distribution :obj:`mpi_array.globale.gndarray`.
        :type converter: :obj:`ToGndarrayConverter`
        :param converter: Used to convert the :obj:`numpy.ndarray` instances
           of :samp:`{func_args}` to :samp:`mpi_array.globale.gndarray` instances.
        :type func: callable
        :param func: Function which computes a new array from the :samp:`*{func_args}`
            arguments and for arguments converted with :samp:`{converter}`.
        :type func_args: sequence of :obj:`numpy.ndarray` or array-like objects
        :param func_args: The arguments for the :samp:`{func}` function.
           Can be comprised of :obj:`numpy.ndarray`, scalars or broadcastable
           sequence (e.g. tuple of scalars) elements.

        .. seealso: :meth:`convert_func_args_to_gndarrays`

        """
        mpi_func_args = self.convert_func_args_to_gndarrays(converter, func_args)
        with func(*mpi_func_args) as mpi_result_ary:
            if mpi_cln_npy_result_ary is None:
                mpi_cln_npy_result_ary = func(*func_args)
            self.compare_results(mpi_cln_npy_result_ary, mpi_result_ary)
        for arg in mpi_func_args:
            if hasattr(arg, "free"):
                arg.free()

    def do_cloned_distribution_test(self, mpi_cln_npy_result_ary, func, *func_args):
        """
        Converts :obj:`numpy.ndarray` elements of :samp:`func_args`
        to :obj:`mpi_array.globale.gndarray` instances distributed as
        the :attr:`mpi_array.comms.DT_CLONED` distribution type.

        :type mpi_cln_npy_result_ary: :obj:`mpi_array.globale.gndarray`
        :param mpi_cln_npy_result_ary: The result returned by :samp:`{func}(*{func_args})`
            converted to a cloned-distribution :obj:`mpi_array.globale.gndarray`.
        :type func: callable
        :param func: Function which computes a new array from the :samp:`*{func_args}`
            arguments.
        :type func_args: sequence of :obj:`numpy.ndarray` or array-like objects
        :param func_args: The arguments for the :samp:`{func}` function.
           Can be comprised of :obj:`numpy.ndarray`, scalars or broadcastable
           sequence (e.g. tuple of scalars) elements.

        .. seealso: :meth:`do_convert_execute_and_compare`, :meth:`compare_results`
        """
        converter = ToGndarrayConverter(locale_type=LT_PROCESS, distrib_type=DT_CLONED)
        self.do_convert_execute_and_compare(mpi_cln_npy_result_ary, converter, func, *func_args)

        converter = ToGndarrayConverter(locale_type=LT_NODE, distrib_type=DT_CLONED, halo=0)
        self.do_convert_execute_and_compare(mpi_cln_npy_result_ary, converter, func, *func_args)

    def do_single_locale_distribution_test(self, mpi_cln_npy_result_ary, func, *func_args):
        """
        Converts :obj:`numpy.ndarray` elements of :samp:`func_args`
        to :obj:`mpi_array.globale.gndarray` instances distributed as
        the :attr:`mpi_array.comms.DT_SINGLE_LOCALE` distribution type.

        :type mpi_cln_npy_result_ary: :obj:`mpi_array.globale.gndarray`
        :param mpi_cln_npy_result_ary: The result returned by :samp:`{func}(*{func_args})`
            converted to a cloned-distribution :obj:`mpi_array.globale.gndarray`.
        :type func: callable
        :param func: Function which computes a new array from the :samp:`*{func_args}`
            arguments.
        :type func_args: sequence of :obj:`numpy.ndarray` or array-like objects
        :param func_args: The arguments for the :samp:`{func}` function.
           Can be comprised of :obj:`numpy.ndarray`, scalars or broadcastable
           sequence (e.g. tuple of scalars) elements.

        .. seealso: :meth:`do_convert_execute_and_compare`
        """
        class Converter(ToGndarrayConverter):

            def __call__(self, np_ary):
                gndary = ToGndarrayConverter.__call__(self, np_ary)
                num_locales = gndary.locale_comms.num_locales
                self.kwargs["inter_locale_rank"] = \
                    ((self.kwargs["inter_locale_rank"] + 1) % num_locales)
                return gndary

        converter = \
            Converter(locale_type=LT_PROCESS, distrib_type=DT_SINGLE_LOCALE, inter_locale_rank=0)
        self.do_convert_execute_and_compare(mpi_cln_npy_result_ary, converter, func, *func_args)

        converter = \
            Converter(locale_type=LT_NODE, distrib_type=DT_SINGLE_LOCALE, inter_locale_rank=0)
        self.do_convert_execute_and_compare(mpi_cln_npy_result_ary, converter, func, *func_args)

    def do_block_distribution_test(self, mpi_cln_npy_result_ary, func, *func_args):
        """
        Converts :obj:`numpy.ndarray` elements of :samp:`func_args`
        to :obj:`mpi_array.globale.gndarray` instances distributed as
        the :attr:`mpi_array.comms.DT_BLOCK` distribution type.

        :type mpi_cln_npy_result_ary: :obj:`mpi_array.globale.gndarray`
        :param mpi_cln_npy_result_ary: The result returned by :samp:`{func}(*{func_args})`
            converted to a cloned-distribution :obj:`mpi_array.globale.gndarray`.
        :type func: callable
        :param func: Function which computes a new array from the :samp:`*{func_args}`
            arguments.
        :type func_args: sequence of :obj:`numpy.ndarray` or array-like objects
        :param func_args: The arguments for the :samp:`{func}` function.
           Can be comprised of :obj:`numpy.ndarray`, scalars or broadcastable
           sequence (e.g. tuple of scalars) elements.

        .. seealso: :meth:`do_convert_execute_and_compare`

        """
        class Converter(ToGndarrayConverter):

            def __init__(self, **kwargs):
                ToGndarrayConverter.__init__(self, **kwargs)
                self.dims = None
                self.axis = 1

            def __call__(self, np_ary):
                if self.dims is None:
                    self.kwargs["dims"] = tuple(_np.zeros((np_ary.ndim,), dtype="int64"))
                gndary = ToGndarrayConverter.__call__(self, np_ary)
                self.axis = min([self.axis, np_ary.ndim])
                self.kwargs["dims"] = _np.ones((np_ary.ndim,), dtype="int64")
                self.kwargs["dims"][self.axis] = 0
                self.kwargs["dims"] = tuple(self.kwargs["dims"])
                self.dims = self.kwargs["dims"]
                self.axis = ((self.axis + 1) % np_ary.ndim)

                return gndary

        converter = Converter(locale_type=LT_PROCESS, distrib_type=DT_BLOCK)
        self.do_convert_execute_and_compare(mpi_cln_npy_result_ary, converter, func, *func_args)

        converter = Converter(locale_type=LT_NODE, distrib_type=DT_BLOCK)
        self.do_convert_execute_and_compare(mpi_cln_npy_result_ary, converter, func, *func_args)

    def do_multi_distribution_tests(self, func, *func_args):
        """
        Compares result of :samp:`{func}` called with :obj:`numpy.ndarray`
        arguments and result of :samp:`{func}` called with :obj:`mpi_array.globale.gndarray`
        arguments.
        Executes :samp:`{func}(*{func_args})` and compares the result
        with :samp:`{func}(*self.convert_func_args_to_gndarrays(converter, {func_args}))`,
        where multiple versions instances of :samp:`converter` are used to generate
        different distributions for the :obj:`mpi_array.globale.gndarray` :samp:`{func}`
        arguments.

        :type func: callable
        :param func: Function which computes a new array from the :samp:`*{func_args}`
            arguments.
        :type func_args: sequence of :obj:`numpy.ndarray` or array-like objects
        :param func_args: The arguments for the :samp:`{func}` function.
           Can be comprised of :obj:`numpy.ndarray`, scalars or broadcastable
           sequence (e.g. tuple of scalars) elements.

        .. seealso: :meth:`do_cloned_distribution_test`, :meth:`do_single_locale_distribution_test`
           , :meth:`do_block_distribution_test` and :meth:`do_convert_execute_and_compare`
        """
        with _asarray(func(*func_args)) as mpi_cln_npy_result_ary:
            self.do_cloned_distribution_test(mpi_cln_npy_result_ary, func, *func_args)
            self.do_single_locale_distribution_test(mpi_cln_npy_result_ary, func, *func_args)
            self.do_block_distribution_test(mpi_cln_npy_result_ary, func, *func_args)

    def test_umath_multiply(self):
        """
        Asserts that binary ufunc multiplication (:obj:`numpy.multiply`) computation
        for :obj:`mpi_array.globale.gndarray` arguments produces same results as
        for :obj:`numpy.ndarray` arguments. Tries various argument combinations
        and different distribution types for the :obj:`mpi_array.globale.gndarray`
        arguments.
        """
        per_axis_size_factor = int(_np.floor(_np.sqrt(float(self.num_node_locales))))
        gshape0 = (41 * per_axis_size_factor + 1, 43 * per_axis_size_factor + 3, 5)
        npy_ary0 = _np.random.uniform(low=0.5, high=1.75, size=gshape0)

        with _asarray(npy_ary0) as cln_ary:
            self.assertTrue(_np.all(npy_ary0 == cln_ary.lndarray_proxy.lndarray))

        def multiply(ary0, ary1):
            return ary0 * ary1

        self.do_multi_distribution_tests(multiply, npy_ary0, 1.0 / 3.0)
        self.do_multi_distribution_tests(multiply, npy_ary0, (0.1, 0.3, 0.5, 0.7, 1.9))
        self.do_multi_distribution_tests(multiply, npy_ary0, npy_ary0)

        gshape1 = gshape0[0:2] + (1,)
        npy_ary1 = _np.random.uniform(low=-0.5, high=2.9, size=gshape1)
        self.do_multi_distribution_tests(multiply, npy_ary0, npy_ary1)
        self.do_multi_distribution_tests(multiply, npy_ary1, npy_ary0)

        npy_ary0 = _np.random.uniform(low=-0.5, high=2.9, size=(gshape0[0], 1, gshape0[2]))
        npy_ary1 = _np.random.uniform(low=-0.5, high=2.9, size=(1, gshape0[1], gshape0[2]))
        self.do_multi_distribution_tests(multiply, npy_ary1, npy_ary0)

    def do_test_umath(self, halo=0, gshape=(32, 48)):
        """
        Test binary op for a :obj:`mpi_array.globale.gndarray` object
        and a scalar.
        """
        with _ones(gshape, dtype="int32", locale_type=_comms.LT_PROCESS, halo=halo) as c:
            # if True:
            #    c = _ones(gshape, dtype="int32", locale_type=_comms.LT_PROCESS, halo=halo)
            c_orig_halo = c.distribution.halo

            self.assertTrue(isinstance(c, _gndarray))
            self.assertTrue((c == 1).all())

            c *= 2
            self.assertTrue((c == 2).all())
            self.assertTrue(_np.all(c.distribution.halo == c_orig_halo))

            with (c + 2) as d:
                self.assertTrue(isinstance(d, _gndarray))
                self.assertEqual(c.dtype, d.dtype)
                self.assertTrue((d == 4).all())
                self.assertTrue(_np.all(d.distribution.halo == c_orig_halo))

    def test_umath_no_halo(self):
        """
        Test binary op for a :obj:`mpi_array.globale.gndarray` object
        and a scalar.
        """
        self.do_test_umath(halo=0)

    def test_umath_halo(self):
        """
        Test binary op for a :obj:`mpi_array.globale.gndarray` object
        and a scalar, test halo is preserved.
        """
        self.do_test_umath(halo=[[1, 2], [3, 4]])

    def do_test_umath_broadcast(self, halo=0, dims=(0, 0, 0)):
        """
        Test binary op for a :obj:`mpi_array.globale.gndarray` objects
        and an *array-like* object which requires requiring broadcast to result shape.
        """
        with \
                _ones(
                    (61, 55, 3),
                    dtype="int32",
                    locale_type=_comms.LT_PROCESS,
                    distrib_type=_comms.DT_BLOCK,
                    dims=dims,
                    halo=halo
                ) as c:
            c_orig_halo = c.distribution.halo

            with (c * (2, 2, 2)) as d:

                self.assertTrue(isinstance(d, _gndarray))
                self.assertEqual(_np.asarray((2, 2, 2)).dtype, d.dtype)
                self.assertSequenceEqual(tuple(c.shape), tuple(d.shape))
                self.assertSequenceEqual(d.distribution.halo.tolist(), c_orig_halo.tolist())
                self.assertTrue((d.view_n == 2).all())
                self.assertTrue((d == 2).all())

    def test_umath_broadcast_no_halo(self):
        """
        Test binary op for a :obj:`mpi_array.globale.gndarray` objects
        and an *array-like* object which requires requiring broadcast to result shape.
        """
        self.do_test_umath_broadcast(halo=0, dims=(0, 0, 0))
        self.do_test_umath_broadcast(halo=0, dims=(1, 1, 0))

    def test_umath_broadcast_halo(self):
        """
        Test binary op for a :obj:`mpi_array.globale.gndarray` objects
        and an *array-like* object which requires requiring broadcast to result shape.
        """
        self.do_test_umath_broadcast(halo=[[1, 2], [3, 4], [2, 1]], dims=(0, 0, 0))
        self.do_test_umath_broadcast(halo=[[1, 2], [3, 4], [2, 1]], dims=(1, 1, 0))

    def do_test_umath_broadcast_upsized_result(
        self,
        halo_a=0,
        halo_b=0,
        dims_a=(0, 0),
        dims_b=(0, 0, 0)
    ):
        """
        Test binary op for two :obj:`mpi_array.globale.gndarray` objects
        with the resulting :obj:`mpi_array.globale.gndarray` object having
        different (larger) shape than that of both inputs.
        """
        with \
                _ones(
                    (19, 3),
                    dtype="int32",
                    locale_type=_comms.LT_PROCESS,
                    distrib_type=_comms.DT_BLOCK,
                    dims=dims_a,
                    halo=halo_a
                ) as a, \
                _ones(
                    (23, 1, 3),
                    dtype="int32",
                    locale_type=_comms.LT_PROCESS,
                    distrib_type=_comms.DT_BLOCK,
                    dims=dims_b,
                    halo=halo_b
                ) as b:

            with (a + b) as d:

                self.assertTrue(isinstance(d, _gndarray))
                self.assertSequenceEqual((b.shape[0], a.shape[0], 3), tuple(d.shape))
                self.assertTrue((d == 2).all())

    def test_umath_broadcast_upsized_result(self):
        """
        Test binary op for two :obj:`mpi_array.globale.gndarray` objects
        with the resulting :obj:`mpi_array.globale.gndarray` object having
        different (larger) shape than that of both inputs.
        """
        self.do_test_umath_broadcast_upsized_result(
            halo_a=0,
            halo_b=0,
            dims_a=(0, 0),
            dims_b=(0, 0, 0)
        )
        self.do_test_umath_broadcast_upsized_result(
            halo_a=0,
            halo_b=0,
            dims_a=(0, 1),
            dims_b=(0, 0, 1)
        )
        self.do_test_umath_broadcast_upsized_result(
            halo_a=[[1, 2], [2, 1]],
            halo_b=0,
            dims_a=(0, 1),
            dims_b=(0, 0, 1)
        )
        self.do_test_umath_broadcast_upsized_result(
            halo_a=0,
            halo_b=[[1, 2], [3, 4], [2, 1]],
            dims_a=(0, 1),
            dims_b=(0, 0, 1)
        )
        self.do_test_umath_broadcast_upsized_result(
            halo_a=[[1, 2], [2, 1]],
            halo_b=[[1, 2], [3, 4], [2, 1]],
            dims_a=(0, 1),
            dims_b=(0, 0, 1)
        )

    def do_test_umath_distributed_broadcast(self, halo_a=0, halo_b=0):
        """
        Test binary op for two :obj:`mpi_array.globale.gndarray` objects
        which requires remote fetch of data when broadcasting to result shape.
        """
        with \
                _ones((61, 53, 5), dtype="int32", locale_type=_comms.LT_PROCESS, halo=halo_a) as a,\
                _ones(a.shape, dtype="int32", locale_type=_comms.LT_PROCESS, halo=halo_b) as b:
            a_orig_halo = a.distribution.halo
            b_orig_halo = b.distribution.halo

            with (a + b) as c:

                self.assertTrue(isinstance(c, _gndarray))
                self.assertTrue((c == 2).all())
                self.assertSequenceEqual(c.distribution.halo.tolist(), a_orig_halo.tolist())

                with \
                        _ones(
                            tuple(a.shape[1:]),
                            dtype=c.dtype,
                            locale_type=_comms.LT_PROCESS,
                            dims=(0, 1),
                            halo=b_orig_halo[1:]
                        ) as twos:

                    twos.fill_h(2)

                    with (a * twos) as d:
                        self.assertTrue(isinstance(d, _gndarray))
                        self.assertSequenceEqual(tuple(a.shape), tuple(d.shape))
                        self.assertTrue((d == 2).all())

    def test_umath_distributed_broadcast_no_halo(self):
        """
        Test binary op for two :obj:`mpi_array.globale.gndarray` objects
        which requires remote fetch of data when broadcasting to result shape.
        """
        self.do_test_umath_distributed_broadcast(halo_a=0, halo_b=0)

    def test_umath_distributed_broadcast_halo(self):
        """
        Test binary op for two :obj:`mpi_array.globale.gndarray` objects
        which requires remote fetch of data when broadcasting to result shape.
        Ghost elements added to arrays.
        """
        self.do_test_umath_distributed_broadcast(halo_a=[[1, 2], [3, 4], [2, 1]], halo_b=0)
        self.do_test_umath_distributed_broadcast(halo_a=0, halo_b=[[1, 2], [3, 4], [2, 1]])
        self.do_test_umath_distributed_broadcast(
            halo_a=[[2, 1], [4, 3], [1, 2]],
            halo_b=[[1, 2], [3, 4], [2, 1]]
        )

    def test_ufunc_casting_arg(self):
        """
        Test ufunc with casting argument.
        """
        self.rank_logger.info("Calling ufunc with 'casting' kwarg.")
        with \
                _ones(
                    (64, 32, 100),
                    dtype="float32",
                    locale_type=_comms.LT_PROCESS,
                    dims=(0, 0, 0),
                    halo=2
                ) as a:

            _np.add(a, 1, casting="same_kind", out=a)
            self.assertTrue((a == 2).all())

    def test_check_equivalent_inter_locale_comms(self):

        with \
                _empty((50, 50, 50), locale_type=LT_NODE) as gary0,\
                _empty((50, 50, 50), locale_type=LT_PROCESS) as gary1:

            if gary1.distribution.num_locales > 1:
                self.assertRaises(
                    ValueError,
                    check_equivalent_inter_locale_comms,
                    (gary0, gary1)
                )

    def test_not_implemented(self):
        uf = _np.add
        with _empty((50, 50, 50), locale_type=LT_NODE) as gary0:

            for method in ["reduce", "accumulate", "reduceat", "at", "outer"]:
                self.assertRaises(
                    TypeError,
                    getattr(uf, method),
                    gary0,
                    5.0
                )

_unittest.main(__name__)

__all__ = [s for s in dir() if not s.startswith('_')]


import time
import pandas as pd
import numpy as np

CHICAGO = 'Chicago'
NYC = 'New York City'
WASHINGTON = 'Washington'
CITY_DATA = { CHICAGO: 'chicago.csv',
              NYC: 'new_york_city.csv',
              WASHINGTON: 'washington.csv' }
days_in_week = [ 'Monday' , 'Tuesday' , 'Wednesday' , 'Thursday' , 'Friday' , 'Saturday' , 'Sunday' ]
months_in_year = [ 'January' , 'February' , 'March' , 'April' , 'May' , 'June' , 'July' , 'August' , 'September' , 'October' , 'November' , 'December' ]
hours = [ '11 PM' , '10 PM' , '9 PM' , '8 PM' , '7 PM' , '6 PM' , '5 PM' , '4 PM' , '3 PM' , '2 PM' , '1 PM' , '12 PM' , '11 AM' , '10 AM' , '9 AM' , '8 AM', '7 AM' , '6 AM' , '5 AM' , '4 AM' , '3 AM' , '2 AM' , '1 AM' , '12 AM' ]

seconds_in_minute = 60
seconds_in_hour = 60 * seconds_in_minute
seconds_in_days = 24 * seconds_in_hour
seconds_in_week = 7 * seconds_in_days

#print(CITY_DATA)
START_TIME = 'Start Time'
END_TIME = 'End Time'
BIRTH_YEAR = 'Birth Year'
START_STATION = 'Start Station'
END_STATION = 'End Station'
TRIP_DURATION = 'Trip Duration'
GENDER = 'Gender'

#colums that are added
START_MONTH = 'Start Month'
START_DAY_OF_WEEK = 'Start Day of Week'

def get_filters():
    """
    Asks user to specify a city, month, and day to analyze.

    Returns:
        (str) city - name of the city to analyze
        (str) month - name of the month to filter by, or "all" to apply no month filter
        (str) day - name of the day of week to filter by, or "all" to apply no day filter
    """
    ALL = 'all'

    print('Hello! Let\'s explore some US bikeshare data!')
    # TO DO: get user input for city (chicago, new york city, washington). HINT: Use a while loop to handle invalid inputs
    while True:
        print("Select the number of the city that you would like to choose ")
        print("1 - Chicago, 2 - New York, 3 - Washington")
        location = input("> ")
        if location == '1' :
            city = CHICAGO
            print("You selected the city Chicago")
            break
        elif location == '2' :
                city = NYC
                print("You selected the city New York")
                break
        elif location == '3' :
                city = WASHINGTON
                print("You selected the city Washington")
                break
        else :
                print("You chose an inavlid option")

    # TO DO: get user input for month (all, january, february, ... , june)
    while True:
        print("Select the number of the month that you want to explore or \"{}\": ".format(ALL))
        print("1 - January .... 6 - June")
        m = input("> ")
        if m == ALL :
            month = None
            break
        try :
            month = int(m)
        except ValueError :
            print("It is an invalid choice")
            continue
        else :
            if month >= 1 and month <= 6 :
                print("You selected " + months_in_year[month - 1])
                break
            elif month <= 12 :
                print("You just have from January to June provided")
                print("It is an invalid choice")
                continue
            else :
                  print("It is an invalid choice")
                  continue
    # get user input for month (all, january, february, ... , june)
    while True:
        print("Select the number of the month you would like to explore or \"{}\": ".format(ALL))
        print("1 - Monday ... 7 - Sunday")
        d = input("> ")
        if d == ALL :
            day = None
            break
        try :
            day = int(d)
        except ValueError :
            print("It is an invalid choice")
            continue
        else :
            if day >= 1 and day<= 7 :
                day -= 1
                print("You selected " + days_in_week[day])
                break
            else :
                print("It is an invalid choice")
                continue

    return city, month, day

def convert_date_time_columns(df) :
                    df[START_TIME] = df[START_TIME].apply(pd.to_datetime)
                    df[END_TIME] = df[END_TIME].apply(pd.to_datetime)
                    df[START_MONTH] = df[START_TIME].dt.month
                    df[START_DAY_OF_WEEK] = df[START_TIME].dt.dayofweek

def load_data(city, month, day):
    """
    Loads data for the specified city and filters by month and day if applicable.

    Args:
        (str) city - name of the city to analyze
        (str) month - name of the month to filter by, or "all" to apply no month filter
        (str) day - name of the day of week to filter by, or "all" to apply no day filter
    Returns:
        df - Pandas DataFrame containing city data filtered by month and day
    """
    #load the data file to the dataframe
    df = pd.read_csv(CITY_DATA[city])

    # convert the start time to datatime
    convert_date_time_columns(df)

    if month is not None :
                    df = df[df[START_MONTH] == month]
    if day is not None :
                    df = df[df[START_DAY_OF_WEEK] == day]

    return df

def time_stats(df, is_display_month, is_display_day_in_week):
    """Displays statistics on the most frequent times of travel.
    Args:
        df - Pandas DataFrame containing city data filtered by month and day
        (bool) is_display_month - if already filtered by month, don't display it
        (bool) is_display_day_in_week - same as above for day of week
    Returns:
        month, day_of_week, hour (tuple of ints) -
            Most frequent times of travel - day_of_week and/or hour may be None if already filtered by these
    """

    print('\nCalculating The Most Frequent Times of Travel...\n')
    start_time = time.time()

    month, day_of_week, hour = None, None, None

    # TO DO: display the most common month
    if is_display_month :
        month = df[START_MONTH].mode()[0]
        print("The most common start month {}".format(months_in_year[month - 1]))

    # TO DO: display the most common day of week
    if is_display_day_in_week :
        day_of_week = df[START_DAY_OF_WEEK].mode()[0]
        print("The most common start day in the whole week {}".format(days_in_week[day_of_week]))

    # TO DO: display the most common start hour
    hour = df[START_TIME].dt.hour.mode()[0]
    print("The most common time {}".format(hours[hour]))
    print("\nThis took %s seconds." % (time.time() - start_time))

    return month, day_of_week, hour

def station_stats(df):
    """Displays statistics on the most popular stations and trip."""

    print('\nCalculating The Most Popular Stations and Trip...\n')
    start_time = time.time()

    # TO DO: display most commonly used start station
    top_start = df['Start Station'].mode()[0]
    common_start_station = (df['Start Station'] == top_start).sum()
    print("The Common Start Station is {} ({} trips)".format(top_start, common_start_station))

    # TO DO: display most commonly used end station
    top_end = df[END_STATION].mode()[0]
    common_end_station = (df[END_STATION] == top_end).sum()
    print('The Common End Station is {} ({} trips)'.format(top_end, common_end_station))

    # TO DO: display most frequent combination of start station and end station trip
    by_start_and_end = df.groupby([START_STATION, END_STATION]).size()
    common_start_end_station = by_start_and_end.idxmax()
    print('The Most Popular Combination Station is \n{}, {} ({} trips)'.format(common_start_end_station[0], common_start_end_station[1], by_start_and_end[common_start_end_station]))

    print("\nThis took %s seconds." % (time.time() - start_time))
    return top_start, top_end, common_start_end_station

def trip_duration_stats(df):
    """Displays statistics on the total and average trip duration"""

    print('\nCalculating Trip Duration...\n')
    start_time = time.time()

    # TO DO: display total travel time
    total_travel_time = df[TRIP_DURATION].sum()
    print('The total travelled time is\n', total_travel_time)

    # TO DO: display mean travel time
    avg_travel_time = df['Trip Duration'].mean()
    print('The average travel time is\n', avg_travel_time)

    print("\nThis took %s seconds." % (time.time() - start_time))
    return total_travel_time, avg_travel_time

def user_stats(df):
    """Displays statistics on bikeshare users.

    Returns:
        user_dict - key type of user, value is count
        gender_dict - same as above, for gender
        earliest_year (int) of birth
        most_recent_year (int) "
        most_common_year (int) "
    """

    print('\nCalculating User Stats...\n')
    start_time = time.time()

    # TO DO: Display counts of user types
    user_type = df.groupby('User Type').count()
    user_dict = {}
    for index, row in user_type.iterrows():
        user_dict[index] = row[0]
        print("{}: {}".format(index, row[0]))

    gender_dict = {}
    if GENDER in df.columns:

    # TO DO: Display counts of gender
        gender = df.groupby(GENDER).count()
        print()
        for index, row in gender.iterrows():
            gender_dict[index] = row[0]
            print("{}: {}".format(index, row[0]))

    the_earliest_year = None
    the_most_recent_year = None
    the_most_common_year = None
    if BIRTH_YEAR in df.columns:

    # TO DO: Display earliest, most recent, and most common year of birth
        birth_year = df[BIRTH_YEAR].dropna()
        the_earliest_year = int(birth_year.min())
        the_most_recent_year = int(birth_year.max())
        the_most_common_year = int(birth_year.mode())

    print("\nThis took %s seconds." % (time.time() - start_time))
    return user_dict, gender_dict, the_earliest_year, the_most_recent_year, the_most_common_year
    print('-'*40)

def main():
    while True:
        city, month, day = get_filters()
        df = load_data(city, month, day)

        time_stats(df, month is None, day is None)
        station_stats(df)
        trip_duration_stats(df)
        user_stats(df)

        restart = input('\nWould you like to restart? Enter yes or no.\n')
        if restart.lower() != 'yes':
            print("Have a great day. Goodbye!")
            break

if __name__ == "__main__":
	main()


"""
Return numpy arrays and Pandas dataframes as LaTeX.

Provides `to_ltx` and `to_clp` which convert numpy arrays and Pandas dataframes
arrays to LaTeX form.
"""

# Note- version must also be set in setup.py
__version__ = '0.82'
__all__ = ['to_clp', 'to_ltx', '__version__']

__author__ = u'Joseph C. Slater'
__license__ = 'MIT'
__copyright__ = 'Copyright 2018 Joseph C. Slater'

import numpy as _np
import pandas as _pd

def to_clp(a, frmt='{:1.2f}', arraytype='bmatrix', imstring='j'):
    r"""
    Return a LaTeX array the the clipboard given a numpy array.

    Parameters
    ----------
    a         : float array
    frmt      : string
        python 3 formatter, optional-
        https://mkaz.tech/python-string-format.html
    arraytype : string
        latex array type- `bmatrix` default, optional
    imstring : string (optional)
        Character for square root of -1. Usually i or j

    Returns
    -------
    out: str
        LaTeX array

    See Also
    --------
    array_to_latex

    Examples
    --------
    >>> import numpy as np
    >>> import array_to_latex as a2l
    >>> A = np.array([[1.23456, 23.45678],[456.23, 8.239521]])
    >>> a2l.to_clp(A, frmt = '{:6.2f}', arraytype = 'array')

    Note that the output is in your clipboard, so you won't see any results.
    See `to_ltx` for further examples.

    """
    b = to_ltx(a, frmt=frmt, arraytype=arraytype, nargout=1, imstring=imstring, print_out=False)
    try:
        import clipboard as _clipboard
        _clipboard.copy(b)
    except ImportError:
        print('\nPackage ''clipboard'' is not installed')
        print('pip install clipboard\nor install via other ',
              'means to use this function')

def _numpyarraytolatex(a, frmt='{:6.2f}', arraytype='bmatrix', nargout=0,
                       imstring='j', row=True, mathform=True):
    r"""Return a LaTeX array given a numpy array.

    Parameters
    ----------
    a         : float array
    frmt      : string
        python 3 formatter, optional-
        https://mkaz.tech/python-string-format.html
    arraytype : string
        latex array type- `bmatrix` default, optional
    imstring : string (optional)
        Character for square root of -1. Usually i or j
    row      : Boolean
        If the array is 1-D, should the output be
            a row (True) or column (False)

    Returns
    -------
    out: str
        LaTeX array

    See Also
    --------
    to_clp

    Examples
    --------
    >>> import numpy as np
    >>> import array_to_latex as a2l
    >>> A = np.array([[1.23456, 23.45678],[456.23, 8.239521]])
    >>> a2l.to_ltx(A, frmt = '{:6.2f}', arraytype = 'array')
    \begin{array}
        1.23 &   23.46\\
      456.23 &    8.24
    \end{array}
    None
    >>> a2l.to_ltx(A, frmt = '{:6.2e}', arraytype = 'array')
    \begin{array}
      1.23e+00 &  2.35e+01\\
      4.56e+02 &  8.24e+00
    \end{array}
    None
    >>> a2l.to_ltx(A, frmt = '{:.3g}', arraytype = 'array')
    \begin{array}
      1.23 &  23.5\\
      456 &  8.24
    \end{array}
    None

    """
    if len(a.shape) > 2:
        raise ValueError('bmatrix can at most display two dimensions')

    if len(a.shape) == 1:
        a = _np.array([a])
        if row is False:
            a = a.T

    out = r'\begin{' + arraytype + '}\n'
    for i in _np.arange(a.shape[0]):
        out = out + ' '
        for j in _np.arange(a.shape[1]):
            if _np.real(a[i, j]) < 0:
                leadstr = ''
            else:
                leadstr = ' '
            if '.' not in frmt.format(a[i, j]):
                dot_space = ' '
            else:
                dot_space = ''
            if _np.iscomplexobj(a[i, j]):
                out = (out + leadstr
                       + math_form(frmt.format(_np.real(a[i, j])),
                                   mathform=mathform)
                       + ' + '
                       + math_form(frmt.format(_np.imag(a[i, j])),
                                   is_imaginary=True,
                                   mathform=mathform)
                       + imstring
                       + dot_space + ' & ')
            else:
                out = (out
                       + leadstr
                       + math_form(frmt.format(_np.real(a[i, j])),
                                   mathform=mathform)
                       + dot_space
                       + r' & ')

        out = out[:-3]
        out = out + '\\\\\n'

    out = out[:-3] + '\n' + r'\end{' + arraytype + '}'

    return out

def _dataframetolatex(df,
                      frmt='{:6.2f}',
                      arraytype='tabular',
                      nargout=0,
                      imstring='j',
                      row=True,
                      mathform=True):
    r"""
    Return a LaTeX array given a Pandas DataFrame array.

    Parameters
    ----------
    a         : float array
    frmt      : string
        python 3 formatter, optional-
        https://mkaz.tech/python-string-format.html
    arraytype : string
        latex array type- `bmatrix` default, optional
    imstring  : string (optional)
        Character for square root of -1. Usually i or j
    row       : Boolean (optional: default True)
        If the array is 1-D, should the output be
            a row (True) or column (False)
    mathform  : Boolean (optional: default True)
        Replace #E# with #\times10^{#}

    Returns
    -------
    out: str
        LaTeX array

    See Also
    --------
    to_clp

    Examples
    --------
    >>> import numpy as np
    >>> import array_to_latex as a2l
    >>> A = np.array([[1.23456, 23.45678],[456.23, 8.239521]])
    >>> a2l.to_ltx(A, frmt = '{:6.2f}', arraytype = 'array')
    \begin{array}
        1.23 &   23.46\\
      456.23 &    8.24
    \end{array}
    None
    >>> a2l.to_ltx(A, frmt = '{:6.2e}', arraytype = 'array')
    \begin{array}
      1.23e+00 &  2.35e+01\\
      4.56e+02 &  8.24e+00
    \end{array}
    None
    >>> a2l.to_ltx(A, frmt = '{:.3g}', arraytype = 'array')
    \begin{array}
      1.23 &  23.5\\
      456 &  8.24
    \end{array}
    None

    """
    columns = df.columns
    rows = df.transpose().columns
    a = _np.array(df)
    out = r'\begin{' + arraytype + '}'

    if arraytype == 'tabular':
        out += r'{l'
        for column in columns:
            out += 'r'
        out += r'}'

    out += '\n'

    if arraytype == 'tabular':
        out += '\\toprule\n'

        out = out + '     '
        for column in columns:
            out += '& ' + column + ' '
        out += r'\\\n'

    if arraytype == 'tabular':
        out += '\\midrule\n'

    for i in _np.arange(a.shape[0]):
        out = out + ' ' + str(rows[i]) + ' & '
        for j in _np.arange(a.shape[1]):
            if isinstance(a[i, j], str):
                leadstr = ' '
                dot_space = (max([len(pet)
                                  for pet in a[:, j]]) - len(a[i, j])) * ' '
                out = (out + leadstr + a[i, j] + dot_space + ' & ')
            else:
                if _np.real(a[i, j]) < 0:
                    leadstr = ''
                else:
                    leadstr = ' '
                if '.' not in frmt.format(a[i, j]):
                    dot_space = ' '
                else:
                    dot_space = ''
                if _np.iscomplexobj(a[i, j]):
                    out = (out + leadstr
                           + math_form(frmt.format(_np.real(a[i, j])),
                                       mathform=mathform)
                           + ' + '
                           + math_form(frmt.format(_np.imag(a[i, j])),
                                       is_imaginary=True,
                                       mathform=mathform)
                           + imstring
                           + dot_space + ' & ')
                else:
                    out = (out + leadstr
                           + math_form(frmt.format(a[i, j]),
                                       mathform=mathform)
                           + dot_space + ' & ')

        out = out[:-3]
        out += '\\\\\n'

    if arraytype == 'tabular':
        out += '\\bottomrule\n'
        out += r'\end{' + arraytype + '}'
    else:
        out = out[:-3] + '\n' + r'\end{' + arraytype + '}'

    return out

def to_ltx(a, frmt='{:1.2f}', arraytype=None, nargout=0,
           imstring='j', row=True, mathform=True, print_out=True):
    r"""
    Print or return a LaTeX array given a numpy array or Pandas dataframe.

    Parameters
    ----------
    a         : float array
    frmt      : string
        python 3 formatter, optional-
        https://mkaz.tech/python-string-format.html
    arraytype : string
        latex array type- `bmatrix` default, optional
    imstring : string (optional)
        Character for square root of -1. Usually i or j
    row        : Boolean (optional: default True)
        If the array is 1-D, should the output be
            a row (True) or column (False)
    mathform  : Boolean (optional: default True)
        Replace #E# with #\times10^{#}

    Returns
    -------
    out: str
        LaTeX array

    See Also
    --------
    to_clp

    Examples
    --------
    >>> import numpy as np
    >>> import array_to_latex as a2l
    >>> A = np.array([[1.23456, 23.45678],[456.23, 8.239521]])
    >>> a2l.to_ltx(A, frmt = '{:6.2f}', arraytype = 'array')
    \begin{array}
        1.23 &   23.46\\
      456.23 &    8.24
    \end{array}
    None
    >>> a2l.to_ltx(A, frmt = '{:6.2e}', arraytype = 'array')
    \begin{array}
      1.23e+00 &  2.35e+01\\
      4.56e+02 &  8.24e+00
    \end{array}
    None
    >>> a2l.to_ltx(A, frmt = '{:.3g}', arraytype = 'array')
    \begin{array}
      1.23 &  23.5\\
      456  &  8.24
    \end{array}
    None

    """
    if isinstance(a, _np.ndarray):

        if arraytype is None:
            arraytype = 'bmatrix'
        latex = _numpyarraytolatex(a, frmt=frmt, arraytype=arraytype,
                                   nargout=nargout, imstring=imstring,
                                   row=row, mathform=mathform)

    elif isinstance(a, _pd.core.frame.DataFrame):

        if arraytype is None:
            arraytype = 'tabular'
        latex = _dataframetolatex(a, frmt=frmt, arraytype=arraytype,
                                  nargout=nargout, imstring=imstring)
    else:
        raise TypeError("Argument should be a "
                        "numpy array or a pandas DataFrame.")
    if print_out is True:
        print(latex)
        return

    return latex

def math_form(number, is_imaginary=False, mathform=True):
    if mathform:
        if 'e' in number:
            number = number.replace('e', '\\times 10^{') + '}'
    return number


"""Utility module for sentiment analysis."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

START_CHAR = 1
END_CHAR = 2
OOV_CHAR = 3

def pad_sentence(sentence, sentence_length):
  """Pad the given sentense at the end.

  If the input is longer than sentence_length,
  the remaining portion is dropped.
  END_CHAR is used for the padding.

  Args:
    sentence: A numpy array of integers.
    sentence_length: The length of the input after the padding.
  Returns:
    A numpy array of integers of the given length.
  """
  sentence = sentence[:sentence_length]
  if len(sentence) < sentence_length:
    sentence = np.pad(sentence, (0, sentence_length - len(sentence)),
                      "constant", constant_values=(START_CHAR, END_CHAR))

  return sentence