from collections import namedtuple

import numpy
import pytest

from thyme.model.fvcom import node_to_centroid
from thyme.model.fvcom import vertical_interpolation

VerticalValues = namedtuple(
    'VerticalValues',
    ['u',
     'v',
     'h',
     'zeta',
     'siglay_centroid',
     'num_nele',
     'num_siglay',
     'time_index',
     'target_depth_default',
     'target_depth_surface',
     'target_depth_deep',
     'expected_u_target_depth_default',
     'expected_v_target_depth_default',
     'expected_u_target_depth_surface',
     'expected_v_target_depth_surface',
     'expected_u_target_depth_deep',
     'expected_v_target_depth_deep'])


@pytest.fixture
def vertical_values():

    time_index = 0
    num_nele = 7
    num_siglay = 20

    target_depth_default = 4.5
    target_depth_surface = 0
    target_depth_deep = 15

    u = numpy.array(
        [[

            [-0.7951693, 0.1987104, 0.20282207, 0.22636837, 0.22338444, 0.22267905, 0.25851583],
            [-0.79512405, 0.17984135, 0.18718982, 0.20845266, 0.20881361, 0.2097686, 0.24099022],
            [-0.79499704, 0.11546519, 0.11235027, 0.12845743, 0.14677778, 0.17036788, 0.14718369],
            [-0.7947852, -0.04345991, -0.0362343, -0.01943639, 0.03145187, 0.05482485, 0.01756552],
            [-0.794485, -0.11175916, -0.10817654, -0.07648413, -0.05387694, -0.02363495, -0.07206099],
            [-0.79409266, -0.15002064, -0.1458673, -0.11296615, -0.09673776, -0.07272214, -0.11500881],
            [-0.7936041, -0.18279588, -0.17216434, -0.15122096, -0.13439055, -0.12422991, -0.15305918],
            [-0.7930144, -0.22257489, -0.21475703, -0.19943172, -0.19317457, -0.18883257, -0.19185503],
            [-0.7923182, -0.26589555, -0.26381484, -0.2563959, -0.2623669, -0.2418433, -0.22976391],
            [-0.7915093, -0.29618508, -0.3025776, -0.30066454, -0.30951524, -0.2785821, -0.2603346],
            [-0.62858844, -0.31844422, -0.3287656, -0.32818496, -0.34010696, -0.2991507, -0.28074494],
            [-0.59509104, -0.3200005, -0.3282557, -0.33371437, -0.3424295, -0.3079574, -0.2808675],
            [-0.5648844, -0.28674564, -0.29246253, -0.3045539, -0.3127165, -0.2805328, -0.26253417],
            [-0.53362864, -0.2308568, -0.23583505, -0.25023454, -0.25649628, -0.2430087, -0.21772045],
            [-0.49928024, -0.1515698, -0.17267239, -0.1813446, -0.19186072, -0.19516239, -0.16393802],
            [-0.45988485, -0.13119768, -0.10612836, -0.11343177, -0.12980227, -0.1302748, -0.12355448],
            [-0.41265842, -0.11103977, -0.09538705, -0.09509068, -0.11300715, -0.10742138, -0.10589793],
            [-0.35249037, -0.09441119, -0.0877323, -0.08854523, -0.10422827, -0.09832109, -0.098252],
            [-0.26714894, -0.08462795, -0.0790456, -0.07945603, -0.09364957, -0.08832049, -0.08927543],
            [-0.10820818, -0.06875627, -0.06421904, -0.06397585, -0.07394486, -0.07045165, -0.07252091]

        ]])

    v = numpy.array(
        [[
            [-0.6447422, 0.04272371, 0.03626542, 0.03629172, 0.02365562, 0.04447182, 0.02940587],
            [-0.64464253, 0.04550566, 0.03864311, 0.03937801, 0.02540983, 0.04566612, 0.03169739],
            [-0.6444702, 0.05464792, 0.05194416, 0.05265207, 0.03223358, 0.04813668, 0.0406685],
            [-0.6442224, 0.05445727, 0.05883322, 0.04983669, 0.0348991, 0.03746793, 0.04016836],
            [-0.64389604, 0.00415152, -0.00149629, -0.00077556, 0.01316911, 0.00495433, 0.01126673],
            [-0.6434879, -0.04056849, -0.0399819, -0.02377908, -0.01824965, -0.0271148, -0.02322751],
            [-0.64299417, -0.05177833, -0.05065423, -0.04175215, -0.03915307, -0.05457062, -0.0499136],
            [-0.6424107, -0.0646683, -0.0626291, -0.05882293, -0.06086947, -0.07759117, -0.06596889],
            [-0.64173275, -0.06379095, -0.06311549, -0.06392502, -0.07486805, -0.08554274, -0.06285658],
            [-0.6409548, -0.05030349, -0.05021876, -0.05602735, -0.06358989, -0.07058521, -0.04798502],
            [-0.48587024, -0.02668872, -0.02443161, -0.03338933, -0.03849523, -0.04692423, -0.03141288],
            [-0.45516193, -0.0020759, 0.00556405, -0.00915105, -0.01343074, -0.02884087, -0.01940906],
            [-0.4284155, 0.0134382, 0.01965301, 0.00118551, -0.00354341, -0.0235792, -0.01681828],
            [-0.40160605, 0.00625094, 0.0087254, -0.00298712, -0.01015626, -0.03078334, -0.02404925],
            [-0.3729973, -0.02794126, -0.02092287, -0.02809138, -0.03298125, -0.04707557, -0.04252515],
            [-0.34107724, -0.02949184, -0.03453156, -0.03266618, -0.04048001, -0.05651411, -0.03880842],
            [-0.30380374, -0.02897151, -0.0317935, -0.02676859, -0.03367185, -0.05099317, -0.03234524],
            [-0.2575078, -0.02654107, -0.02930111, -0.0243634, -0.03028261, -0.04743995, -0.02939028],
            [-0.19346614, -0.02417713, -0.0263748, -0.02134923, -0.02650077, -0.04303462, -0.02615478],
            [-0.07716001, -0.01982844, -0.02149794, -0.01682694, -0.02036272, -0.03471823, -0.02087618]

        ]])

    h = numpy.array(
        [

            1.372000018755595, 1.91399347, 4.6028045, 3.29994912, 30.32372402, 9.3237240, 12.19253056

        ])

    zeta = numpy.array(
        [
            -1.1558676163355508, -0.1443424, -0.14408446, -0.15902775, -0.16805193, -0.16938841, -0.16457143,

        ])

    siglay_centroid = numpy.array(
        [

            [-0.025, -0.02500004, -0.02500004, -0.02500004, -0.02500004, -0.02500004, -0.02500004],
            [-0.075, -0.07500005, -0.07500005, -0.07500005, -0.07500005, -0.07500005, -0.07500005],
            [-0.125, -0.12500003, -0.12500003, -0.12500003, -0.12500003, -0.12500003, -0.12500003],
            [-0.17500001, -0.17500001, -0.17500001, -0.17500001, -0.17500001, -0.17500001, -0.17500001],
            [-0.225, -0.22500002, -0.22500002, -0.22500002, -0.22500002, -0.22500002, -0.22500002],
            [-0.275, -0.27500007, -0.27500007, -0.27500007, -0.27500007, -0.27500007, -0.27500007],
            [-0.325, -0.32500005, -0.32500005, -0.32500005, -0.32500005, -0.32500005, -0.32500005],
            [-0.375, -0.37500003, -0.37500003, -0.37500003, -0.37500003, -0.37500003, -0.37500003],
            [-0.425, -0.42500004, -0.42500004, -0.42500004, -0.42500004, -0.42500004, -0.42500004],
            [-0.475, -0.47500002, -0.47500002, -0.47500002, -0.47500002, -0.47500002, -0.47500002],
            [-0.525, -0.52500004, -0.52500004, -0.52500004, -0.52500004, -0.52500004, -0.52500004],
            [-0.57500005, -0.57500005, -0.57500005, -0.57500005, -0.57500005, -0.57500005, -0.57500005],
            [-0.625, -0.6250001, -0.6250001, -0.6250001, -0.6250001, -0.6250001, -0.6250001],
            [-0.67499995, -0.6750001, -0.6750001, -0.6750001, -0.6750001, -0.6750001, -0.6750001],
            [-0.725, -0.725, -0.725, -0.725, -0.725, -0.725, -0.725],
            [-0.775, -0.7750001, -0.7750001, -0.7750001, -0.7750001, -0.7750001, -0.7750001],
            [-0.82500005, -0.82500005, -0.82500005, -0.82500005, -0.82500005, -0.82500005,  -0.82500005],
            [-0.875, -0.87500006, -0.87500006, -0.87500006, -0.87500006, -0.87500006, -0.87500006],
            [-0.92499995, -0.9250001, -0.9250001, -0.9250001, -0.9250001, -0.9250001,  -0.9250001],
            [-0.975, -0.975, -0.975, -0.975, -0.975, -0.975, -0.975]

        ])

    expected_u_target_depth_default = numpy.array(
        [
            -0.7100488057968346, -0.30731464, -0.31567158, -0.31442473, 0.09090091, -0.28539867, -0.19117865
        ])

    expected_v_target_depth_default = numpy.array(
        [
            -0.5634124710903746, -0.03849612, -0.0373252, -0.04470835,  0.03352506, -0.06274381, -0.06568897
        ])

    expected_u_target_depth_surface = numpy.array(
        [
            -0.79519192, 0.20814494, 0.21063821, 0.23532624, 0.23066987, 0.22913428, 0.26727865
        ])

    expected_v_target_depth_surface = numpy.array(
        [
            -0.64479204, 0.04133273, 0.03507657, 0.03474857, 0.02277851, 0.04387467, 0.02826011
        ])

    expected_u_target_depth_deep = numpy.array(
        [
            -0.71004887, -0.30731464, -0.31567158, -0.31442473, -0.32323185, -0.28886639, -0.27053976
        ])

    expected_v_target_depth_deep = numpy.array(
        [
            -0.56341252, -0.03849612, -0.0373252, -0.04470835, -0.05233803, -0.05875473, -0.03969896
        ])

    return VerticalValues(u, v, h, zeta, siglay_centroid, num_nele, num_siglay, time_index, target_depth_default,
                          target_depth_surface, target_depth_deep, expected_u_target_depth_default,
                          expected_v_target_depth_default, expected_u_target_depth_surface,
                          expected_v_target_depth_surface, expected_u_target_depth_deep, expected_v_target_depth_deep)


def test_vertical_interpolation(vertical_values):
    """Test vertical interpolation."""
    u_target_depth, v_target_depth = vertical_interpolation(vertical_values.u, vertical_values.v, vertical_values.h,
                                                            vertical_values.zeta, vertical_values.siglay_centroid,
                                                            vertical_values.num_nele, vertical_values.num_siglay,
                                                            vertical_values.time_index,
                                                            vertical_values.target_depth_default)

    # print(f"u_target_depth: {u_target_depth}")
    # print(f"v_target_depth: {v_target_depth}")
    assert numpy.allclose(u_target_depth, vertical_values.expected_u_target_depth_default)
    assert numpy.allclose(v_target_depth, vertical_values.expected_v_target_depth_default)


def test_vertical_interpolation_at_surface(vertical_values):
    """Test vertical interpolation."""
    u_target_depth, v_target_depth = vertical_interpolation(vertical_values.u, vertical_values.v, vertical_values.h,
                                                            vertical_values.zeta, vertical_values.siglay_centroid,
                                                            vertical_values.num_nele, vertical_values.num_siglay,
                                                            vertical_values.time_index,
                                                            vertical_values.target_depth_surface)

    # print(f"u_target_depth_surface: {u_target_depth}")
    # print(f"v_target_depth_surface: {v_target_depth}")
    assert numpy.allclose(u_target_depth, vertical_values.expected_u_target_depth_surface)
    assert numpy.allclose(v_target_depth, vertical_values.expected_v_target_depth_surface)


def test_vertical_interpolation_deep(vertical_values):
    """Test vertical interpolation."""
    u_target_depth, v_target_depth = vertical_interpolation(vertical_values.u, vertical_values.v, vertical_values.h,
                                                            vertical_values.zeta, vertical_values.siglay_centroid,
                                                            vertical_values.num_nele, vertical_values.num_siglay,
                                                            vertical_values.time_index,
                                                            vertical_values.target_depth_deep)

    # print(f"u_target_depth_deep: {u_target_depth}")
    # print(f"v_target_depth_deep: {v_target_depth}")
    assert numpy.allclose(u_target_depth, vertical_values.expected_u_target_depth_deep)
    assert numpy.allclose(v_target_depth, vertical_values.expected_v_target_depth_deep)


NodeToCentroidValues = namedtuple(
    'NodeToCentroidValues',
    ['zeta',
     'h',
     'lat_node',
     'lon_node',
     'lat_centroid',
     'lon_centroid',
     'time_index',
     'expected_h_centroid',
     'expected_zeta_centroid'])


@pytest.fixture
def node_to_centroid_values():
    zeta = numpy.array(
        [
            [-0.146551, -0.146383, -0.145782, -0.146226, -0.145688, -0.145071, -0.146047, -0.145581, -0.145099,
             -0.144439, -0.145729, -0.145431, -0.144983, -0.144627, -0.144069, -0.145415, -0.145071, -0.144579,
             -0.144391, -0.143806, -0.144057,

             -0.169269, -0.169959, -0.170395, -0.170738, -0.166229, -0.166729, -0.167032, -0.162118, -0.162876,
             -0.163448, -0.163812, -0.158493, -0.159403, -0.159949, -0.160389
             ]
        ])

    h = numpy.array(
        [
            1.6438, 1.8724, 1.875, 2.0302, 1.9946, 2.0446, 2.1169, 2.1307, 2.1457, 2.2134, 2.1768, 2.178, 2.2341,
            2.2623, 2.331, 2.1949, 2.2047, 2.2817, 2.3537, 2.434, 2.404,

            31.461, 31.874, 32.37, 32.361, 30.565, 31.089, 31.196, 29.231, 29.033, 28.687, 28.761, 26.053, 25.788,
            26.267, 26.835
        ])
    lat_centroid = numpy.array(
        [
            30.260946, 30.262287, 30.260996, 30.259796, 30.26362, 30.262392, 30.261309, 30.260002, 30.258938, 30.264635,
            30.263693, 30.262672, 30.261517, 30.260561, 30.259289, 30.258574, 30.264799, 30.264191, 30.263008,
            30.262444, 30.261324, 30.260527, 30.259325, 30.259058, 30.260839, 30.259773,

            29.969425, 29.963093, 29.967611, 29.961212, 29.96579, 29.959322, 29.980326, 29.985025,
            29.978645, 29.983408, 29.976954, 29.981771, 29.996161, 30.001049, 29.994678, 29.999624,
            29.993172, 29.998173
        ])

    lon_centroid = numpy.array(
        [
            -88.126221, -88.124329, -88.124817, -88.123871, -88.122406, -88.122925, -88.12204, -88.122498, -88.12149,
            -88.120361, -88.120972, -88.120026, -88.120575, -88.119843, -88.1203, -88.119263, -88.118958, -88.118195,
            -88.118652, -88.117645, -88.117676, -88.118591, -88.118286, -88.117096, -88.116516, -88.116272,

            -88.264221, -88.25528, -88.245392, -88.236328, -88.226379, -88.217194, -88.265015,
            -88.256958, -88.246277, -88.238129, -88.227295, -88.219116, -88.260498, -88.253448,
            -88.241791, -88.23468, -88.222839, -88.215698
        ])

    lat_node = numpy.array(
        [
            30.260969, 30.262211, 30.259659, 30.263531, 30.26112, 30.258608, 30.264805, 30.262526, 30.260281, 30.257925,
            30.265354, 30.263748, 30.261744, 30.259661, 30.258137, 30.265295, 30.263531, 30.262054, 30.260176, 30.25886,
            30.260286,

            29.958637, 29.956696, 29.954746, 29.952793, 29.973944, 29.972195, 29.970428, 29.991339, 29.989796,
            29.988235, 29.986652, 30.007345, 30.006006, 30.004629, 30.003235
        ])

    lon_node = numpy.array(
        [
            -88.127625, -88.125763, -88.125275, -88.12384, -88.123413, -88.122894, -88.121857, -88.12149, -88.121185,
            -88.120422, -88.119629, -88.119568, -88.119049, -88.119293, -88.118073, -88.117676, -88.11734, -88.116516,
            -88.117462, -88.115784, -88.11557,

            -88.265167, -88.246277, -88.227234, -88.207947, -88.254395, -88.235504, -88.2164, -88.267578,
            -88.248901, -88.230011, -88.210937, -88.265045, -88.246429, -88.2276, -88.208527
        ])

    time_index = 0

    expected_h_centroid = numpy.array(
        [
            1.79706003, 1.96634453, 1.91399347, 1.97068808, 2.09215354, 2.05125885, 2.08959921, 2.06159611, 2.13525152,
            2.15697391, 2.14185806, 2.1812977, 2.17015265, 2.2140166,  2.20713333, 2.26888473, 2.18323025, 2.19253056,
            2.20559465, 2.23989205, 2.28985895, 2.28394075, 2.31515434, 2.37337259, 2.3464489, 2.39725858,

            30.74697778, 31.29994912, 31.1760438, 31.77788929, 31.55164158, 31.9756924, 30.00374208, 29.60978718,
            30.22881187, 29.6028045, 30.32372402, 29.54809818, 28.10520577, 26.95770733, 27.83656214, 26.91390017,
            27.90491454, 27.28733678

        ])

    expected_zeta_centroid = numpy.array(
        [
           -0.1462386, -0.14609691, -0.14595087, -0.1455159, -0.14595329, -0.14583367, -0.14545806, -0.14528592,
           -0.14486764, -0.14573735, -0.1456864, -0.14532972, -0.14522106, -0.14490293, -0.14472167, -0.14437832,
           -0.14552505, -0.14530566, -0.14516168, -0.14487977, -0.14465093, -0.14466511, -0.1443643, -0.14408686,
           -0.1443424, -0.14408446,

           -0.16695492, -0.16848552, -0.16763914, -0.16902775, -0.16805193, -0.16938841, -0.16457143, -0.16374129,
           -0.16527793, -0.16435127, -0.16573603, -0.16476413, -0.16116232, -0.16025768, -0.16190908, -0.16093319,
           -0.16240301, -0.16138305
        ])

    return NodeToCentroidValues(zeta, h, lat_node, lon_node, lat_centroid, lon_centroid, time_index,
                                expected_h_centroid, expected_zeta_centroid)


def test_node_to_centroid(node_to_centroid_values):
    """Test horizontal interpolation from nodes to centroids(elements)."""
    h_centroid, zeta_centroid = node_to_centroid(node_to_centroid_values.zeta, node_to_centroid_values.h,
                                                 node_to_centroid_values.lon_node, node_to_centroid_values.lat_node,
                                                 node_to_centroid_values.lon_centroid,
                                                 node_to_centroid_values.lat_centroid,
                                                 node_to_centroid_values.time_index)

    # print(f"h_centroid: {h_centroid}")
    # print(f"zeta_centroid: {zeta_centroid}")
    assert numpy.allclose(h_centroid, node_to_centroid_values.expected_h_centroid)
    assert numpy.allclose(zeta_centroid, node_to_centroid_values.expected_zeta_centroid)
