#  Copyright (c) 2021 Mira Geoscience Ltd.
#
#  This file is part of geoapps.
#
#  geoapps is distributed under the terms and conditions of the MIT License
#  (see LICENSE file at the root of this source code package).

import itertools

import numpy as np
from discretize import TreeMesh
from geoh5py.objects import Octree
from geoh5py.workspace import Workspace

from geoapps.utils.utils import (
    octree_2_treemesh,
    rotate_xy,
    running_mean,
    treemesh_2_octree,
    weighted_average,
)


def test_rotation_xy():
    vec = np.c_[1, 0, 0]
    rot_vec = rotate_xy(vec, [0, 0], 45)

    assert (
        np.linalg.norm(np.cross(rot_vec, [0.7071, 0.7071, 0])) < 1e-8
    ), "Error on positive rotation about origin."

    rot_vec = rotate_xy(vec, [1, 1], -90)

    assert (
        np.linalg.norm(np.cross(rot_vec, [0, 1, 0])) < 1e-8
    ), "Error on negative rotation about point."


def test_running_mean():
    vec = np.random.randn(100)
    mean_forw = running_mean(vec, method="forward")
    mean_back = running_mean(vec, method="backward")
    mean_cent = running_mean(vec, method="centered")

    mean_test = (vec[1:] + vec[:-1]) / 2

    assert (
        np.linalg.norm(mean_back[:-1] - mean_test) < 1e-12
    ), "Backward averaging does not match expected values."
    assert (
        np.linalg.norm(mean_forw[1:] - mean_test) < 1e-12
    ), "Forward averaging does not match expected values."
    assert (
        np.linalg.norm((mean_test[1:] + mean_test[:-1]) / 2 - mean_cent[1:-1]) < 1e-12
    ), "Centered averaging does not match expected values."


def test_weigted_average():

    # in loc == out loc -> in val == out val
    xyz_out = np.array([[0, 0, 0]])
    xyz_in = np.array([[0, 0, 0]])
    values = [np.array([99])]
    out = weighted_average(xyz_in, xyz_out, values)
    assert out[0] == 99

    # two point same distance away -> arithmetic mean
    xyz_in = np.array([[1, 0, 0], [0, 1, 0]])
    values = [np.array([99, 100])]
    out = weighted_average(xyz_in, xyz_out, values)
    assert (out[0] - 99.5) < 1e-10

    # two points different distance away but close to infinity -> arithmetic mean
    xyz_in = np.array([[1e30, 0, 0], [1e30 + 1, 0, 0]])
    values = [np.array([99, 100])]
    out = weighted_average(xyz_in, xyz_out, values)
    assert (out[0] - 99.5) < 1e-10

    # one point close to infinity, one not -> out val is near-field value
    xyz_in = np.array([[1, 0, 0], [1e30, 0, 0]])
    values = [np.array([99, 100])]
    out = weighted_average(xyz_in, xyz_out, values)
    assert (out[0] - 99.0) < 1e-10

    # one values vector and n out locs -> one out vector of length 20
    xyz_in = np.random.rand(10, 3)
    xyz_out = np.random.rand(20, 3)
    values = [np.random.rand(10)]
    out = weighted_average(xyz_in, xyz_out, values)
    assert len(out) == 1
    assert len(out[0]) == 20

    # two values vectors and n out locs -> two out vectors of length 20 each
    xyz_in = np.random.rand(10, 3)
    xyz_out = np.random.rand(20, 3)
    values = [np.random.rand(10), np.random.rand(10)]
    out = weighted_average(xyz_in, xyz_out, values)
    assert len(out) == 2
    assert len(out[0]) == 20
    assert len(out[1]) == 20

    # max distance keeps out points from average
    xyz_in = np.array([[1, 0, 0], [3, 0, 0]])
    xyz_out = np.array([[0, 0, 0]])
    values = [np.array([1, 100])]
    out = weighted_average(xyz_in, xyz_out, values, max_distance=2)
    assert out[0] == 1

    # n caps the number of points that go into the average
    xyz_in = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
    xyz_out = np.array([[0, 0, 0]])
    values = [np.array([1, 2, 3])]
    out = weighted_average(xyz_in, xyz_out, values, n=3)
    assert out[0] == 2
    out = weighted_average(xyz_in, xyz_out, values, n=2)
    assert out[0] == 1.5

    # return indices with n=1 returns closest in loc to the out loc
    xyz_in = np.array([[2, 0, 0], [0, 1, 0], [0, 0, 2]])
    xyz_out = np.array([[0, 0, 0]])
    values = [np.array([1, 2, 3])]
    out, ind = weighted_average(xyz_in, xyz_out, values, n=1, return_indices=True)
    assert ind[0][0] == 1

    # threshold >> r -> arithmetic mean
    xyz_in = np.array([[1, 0, 0], [0, 100, 0], [0, 0, 1000]])
    xyz_out = np.array([[0, 0, 0]])
    values = [np.array([1, 2, 3])]
    out = weighted_average(xyz_in, xyz_out, values, threshold=1e30)
    assert out[0] == 2


def test_treemesh_2_octree():

    ws = Workspace("./FlinFlon.geoh5")
    mesh = TreeMesh([[10] * 16, [10] * 4, [10] * 8], [0, 0, 0])
    mesh.insert_cells([10, 10, 10], mesh.max_level, finalize=True)
    omesh = treemesh_2_octree(ws, mesh, name="test_mesh")
    assert omesh.n_cells == mesh.n_cells
    assert np.all((omesh.centroids - mesh.cell_centers[mesh._ubc_order]) < 1e-14)
    expected_refined_cells = [
        (0, 0, 6),
        (0, 0, 7),
        (1, 0, 6),
        (1, 0, 7),
        (0, 1, 6),
        (0, 1, 7),
        (1, 1, 6),
        (1, 1, 7),
    ]
    ijk_refined = omesh.octree_cells[["I", "J", "K"]][
        omesh.octree_cells["NCells"] == 1
    ].tolist()
    assert np.all([k in ijk_refined for k in expected_refined_cells])
    assert np.all([k in expected_refined_cells for k in ijk_refined])


def test_octree_2_treemesh():
    ws = Workspace("./FlinFlon.geoh5")
    mesh = TreeMesh([[10] * 4, [10] * 4, [10] * 4], [0, 0, 0])
    mesh.insert_cells([5, 5, 5], mesh.max_level, finalize=True)
    omesh = treemesh_2_octree(ws, mesh)
    for p in itertools.product("uvw", repeat=3):
        omesh.origin = [0, 0, 0]
        for axis in "uvw":
            attr = axis + "_cell_size"
            setattr(omesh, attr, np.abs(getattr(omesh, attr)))
        for axis in np.unique(p):
            attr = axis + "_cell_size"
            setattr(omesh, attr, -1 * getattr(omesh, attr))
            omesh.origin["xyz"["uvw".find(axis)]] = 40
        tmesh = octree_2_treemesh(omesh)
        assert np.all((tmesh.cell_centers - mesh.cell_centers) < 1e-14)
