# resqpy test module including well blocking of trajectories against challenging grid geometries

version = '29th April 2021'

import logging

log = logging.getLogger(__name__)

import math as maths
import os

import numpy as np
import pandas as pd

import resqpy.crs as rqc
import resqpy.grid as grr
import resqpy.model as rq
import resqpy.olio.uuid as bu
import resqpy.olio.vector_utilities as vec
import resqpy.olio.xml_et as rqet
import resqpy.rq_import as rqi
import resqpy.well as rqw


def test_s_bend_fn(tmp_path, epc = None):

    if epc is None:
        # use pytest temporary directory fixture
        # https://docs.pytest.org/en/stable/tmpdir.html
        epc = str(os.path.join(tmp_path, f"{bu.new_uuid()}.epc"))

    # create s-bend grid

    nk = 5
    nj = 12
    ni_tail = 5
    ni_bend = 18
    ni_half_mid = 2
    ni = 2 * (ni_tail + ni_bend + ni_half_mid)

    total_thickness = 12.0
    layer_thickness = total_thickness / float(nk)
    flat_dx_di = 10.0
    horst_dx_dk = 0.25 * flat_dx_di / float(nk)
    horst_dz = 1.73 * layer_thickness
    horst_half_dx = horst_dx_dk * horst_dz / layer_thickness
    dy_dj = 8.0
    top_depth = 100.0

    assert ni_bend % 2 == 0, 'ni_bend must be even for horizontal faulting'
    assert ni_tail >= 5, 'ni_tail must be at least 5 for horst blocks'

    bend_theta_di = maths.pi / float(ni_bend)
    outer_radius = 2.0 * total_thickness

    bend_a_centre_xz = (flat_dx_di * float(ni_tail), top_depth + outer_radius)
    bend_b_centre_xz = (flat_dx_di * float(ni_tail - 2.0 * ni_half_mid),
                        top_depth + 3.0 * outer_radius - total_thickness)

    points = np.empty((nk + 1, nj + 1, ni + 1, 3))

    for k in range(nk + 1):
        if k == nk // 2 + 1:
            points[k] = points[k - 1]  # pinched out layer
        else:
            for i in range(ni + 1):
                if i < ni_tail + 1:
                    x = flat_dx_di * float(i)
                    z = top_depth + float(k) * layer_thickness  # will introduce a thick layer after pinchout
                elif i < ni_tail + ni_bend:
                    theta = (i - ni_tail) * bend_theta_di
                    radius = outer_radius - float(k) * layer_thickness
                    x = bend_a_centre_xz[0] + radius * maths.sin(theta)
                    z = bend_a_centre_xz[1] - radius * maths.cos(theta)
                elif i < ni_tail + ni_bend + 2 * ni_half_mid + 1:
                    x = flat_dx_di * float(ni_tail - (i - (ni_tail + ni_bend)))
                    z = top_depth + 2.0 * outer_radius - float(k) * layer_thickness
                elif i < ni_tail + 2 * ni_bend + 2 * ni_half_mid:
                    theta = (i - (ni_tail + ni_bend + 2 * ni_half_mid)) * bend_theta_di
                    radius = outer_radius - float(nk - k) * layer_thickness
                    x = bend_b_centre_xz[0] - radius * maths.sin(theta)
                    z = bend_b_centre_xz[1] - radius * maths.cos(theta)
                else:
                    x = flat_dx_di * float((i - (ni - ni_tail)) + ni_tail - 2 * ni_half_mid)
                    if i == ni - 1 or i == ni - 4:
                        x += horst_dx_dk * float(k)
                    elif i == ni - 2 or i == ni - 3:
                        x -= horst_dx_dk * float(k)
                    z = top_depth + 4.0 * outer_radius + float(k) * layer_thickness - 2.0 * total_thickness
                points[k, :, i] = (x, 0.0, z)

    for j in range(nj + 1):
        points[:, j, :, 1] = dy_dj * float(j)

    model = rq.Model(epc_file = epc, new_epc = True, create_basics = True, create_hdf5_ext = True)

    grid = grr.Grid(model)

    crs = rqc.Crs(model)
    crs_node = crs.create_xml()
    if model.crs_root is None:
        model.crs_root = crs_node

    grid.grid_representation = 'IjkGrid'
    grid.extent_kji = np.array((nk, nj, ni), dtype = 'int')
    grid.nk, grid.nj, grid.ni = nk, nj, ni
    grid.k_direction_is_down = True  # dominant layer direction, or paleo-direction
    grid.pillar_shape = 'straight'
    grid.has_split_coordinate_lines = False
    grid.k_gaps = None
    grid.crs_uuid = crs.uuid
    grid.crs_root = crs_node

    grid.points_cached = points

    grid.geometry_defined_for_all_pillars_cached = True
    grid.geometry_defined_for_all_cells_cached = True
    grid.grid_is_right_handed = crs.is_right_handed_xyz()

    grid.write_hdf5_from_caches()
    grid.create_xml()

    # create a well trajectory and md datum

    def df_trajectory(x, y, z):
        N = len(x)
        assert len(y) == N and len(z) == N
        df = pd.DataFrame(columns = ['MD', 'X', 'Y', 'Z'])
        md = np.zeros(N)
        for n in range(N - 1):
            md[n + 1] = md[n] + vec.naive_length((x[n + 1] - x[n], y[n + 1] - y[n], z[n + 1] - z[n]))
        df.MD = md
        df.X = x
        df.Y = y
        df.Z = z
        return df

    x = np.array(
        [0.0, flat_dx_di * float(ni_tail) + outer_radius, flat_dx_di * (float(ni_tail) - 0.5), 0.0, -outer_radius])
    y = np.array([0.0, dy_dj * 0.5, dy_dj * float(nj) / 2.0, dy_dj * (float(nj) - 0.5), dy_dj * float(nj)])
    z = np.array([
        0.0, top_depth - total_thickness, top_depth + 2.0 * outer_radius - total_thickness / 2.0,
        top_depth + 3.0 * outer_radius - total_thickness, top_depth + 4.0 * outer_radius
    ])

    df = df_trajectory(x, y, z)

    datum = rqw.MdDatum(model, crs_uuid = crs.uuid, location = (x[0], y[0], z[0]))
    datum.create_xml()

    trajectory = rqw.Trajectory(model, md_datum = datum, data_frame = df, length_uom = 'm', well_name = 'ANGLED_WELL')

    rqc.Crs(model, uuid = rqet.uuid_for_part_root(trajectory.crs_root)).uuid == crs.uuid

    trajectory.write_hdf5()
    trajectory.create_xml()

    # add more wells

    x = np.array([
        0.0, flat_dx_di * float(ni_tail), flat_dx_di * 2.0 * float(ni_tail - ni_half_mid) + outer_radius, -outer_radius
    ])
    y = np.array([0.0, dy_dj * float(nj) * 0.59, dy_dj * 0.67, dy_dj * 0.5])
    z = np.array([
        0.0, top_depth - total_thickness, top_depth + 4.0 * outer_radius - 1.7 * total_thickness,
        top_depth + 4.0 * outer_radius - 1.7 * total_thickness
    ])

    df = df_trajectory(x, y, z)

    traj_2 = rqw.Trajectory(model, md_datum = datum, data_frame = df, length_uom = 'm', well_name = 'HORST_WELL')
    traj_2.write_hdf5()
    traj_2.create_xml()
    traj_2.control_points

    x = np.array([0.0, 0.0, 0.0])
    y = np.array([0.0, dy_dj * float(nj) * 0.53, dy_dj * float(nj) * 0.53])
    z = np.array([0.0, top_depth - total_thickness, top_depth + 4.0 * outer_radius])

    df = df_trajectory(x, y, z)

    traj_3 = rqw.Trajectory(model, md_datum = datum, data_frame = df, length_uom = 'm', well_name = 'VERTICAL_WELL')
    traj_3.write_hdf5()
    traj_3.create_xml()
    traj_3.control_points

    n_x = flat_dx_di * float(ni_tail) * 0.48
    n_y = dy_dj * float(nj) / 9.1
    o_y = -dy_dj * 0.45
    nd_x = n_y / 3.0
    x = np.array([
        0.0, n_x, n_x, n_x + nd_x, n_x + 2.0 * nd_x, n_x + 3.0 * nd_x, n_x + 4.0 * nd_x, n_x + 5.0 * nd_x,
        n_x + 6.0 * nd_x, n_x + 7.0 * nd_x, n_x + 8.0 * nd_x, n_x + 8.0 * nd_x
    ])
    y = np.array([
        0.0, o_y, o_y + n_y, o_y + 2.0 * n_y, o_y + 3.0 * n_y, o_y + 4.0 * n_y, o_y + 5.0 * n_y, o_y + 6.0 * n_y,
        o_y + 7.0 * n_y, o_y + 8.0 * n_y, o_y + 9.0 * n_y, o_y + 10.0 * n_y
    ])
    n_z1 = top_depth + total_thickness * 0.82
    n_z2 = top_depth - total_thickness * 0.17
    z = np.array([0.0, n_z1, n_z1, n_z2, n_z2, n_z1, n_z1, n_z2, n_z2, n_z1, n_z1, n_z2])

    df = df_trajectory(x, y, z)

    traj_4 = rqw.Trajectory(model, md_datum = datum, data_frame = df, length_uom = 'm', well_name = 'NESSIE_WELL')
    traj_4.write_hdf5()
    traj_4.create_xml()
    traj_4.control_points

    # block wells against grid geometry

    log.info('unfaulted grid blocking of well ' + str(rqw.well_name(trajectory)))
    bw = rqw.BlockedWell(model, grid = grid, trajectory = trajectory)
    bw.write_hdf5()
    bw.create_xml()
    assert bw.cell_count == 19

    log.info('unfaulted grid blocking of well ' + str(rqw.well_name(traj_2)))
    bw_2 = rqw.BlockedWell(model, grid = grid, trajectory = traj_2)
    bw_2.write_hdf5()
    bw_2.create_xml()
    assert bw_2.cell_count == 33

    log.info('unfaulted grid blocking of well ' + str(rqw.well_name(traj_3)))
    bw_3 = rqw.BlockedWell(model, grid = grid, trajectory = traj_3)
    bw_3.write_hdf5()
    bw_3.create_xml()
    assert bw_3.cell_count == 18

    log.info('unfaulted grid blocking of well ' + str(rqw.well_name(traj_4)))
    bw_4 = rqw.BlockedWell(model, grid = grid, trajectory = traj_4)
    bw_4.write_hdf5()
    bw_4.create_xml()
    assert bw_4.cell_count == 26

    # derive a faulted version of the grid

    cp = grid.corner_points(cache_cp_array = True).copy()

    # IK plane faults
    cp[:, 3:, :, :, :, :, :] += (flat_dx_di * 0.7, 0.0, layer_thickness * 1.3)
    cp[:, 5:, :, :, :, :, :] += (flat_dx_di * 0.4, 0.0, layer_thickness * 0.9)
    cp[:, 8:, :, :, :, :, :] += (flat_dx_di * 0.3, 0.0, layer_thickness * 0.6)

    # JK plane faults
    cp[:, :, ni_tail + ni_bend // 2:, :, :, :, 0] -= flat_dx_di * 0.57  # horizontal break mid top bend
    cp[:, :, ni_tail + ni_bend + ni_half_mid:, :, :, :, 2] += layer_thickness * 1.27  # vertical break in mid section

    # zig-zag fault
    j_step = nj // (ni_tail - 2)
    for i in range(ni_tail - 1):
        j_start = i * j_step
        if j_start >= nj:
            break
        cp[:, j_start:, i, :, :, :, 2] += 1.1 * total_thickness

    # JK horst blocks
    cp[:, :, ni - 4, :, :, :, :] -= (horst_half_dx, 0.0, horst_dz)
    cp[:, :, ni - 3:, :, :, :, 0] -= 2.0 * horst_half_dx
    cp[:, :, ni - 2, :, :, :, :] += (-horst_half_dx, 0.0, horst_dz)
    cp[:, :, ni - 1:, :, :, :, 0] -= 2.0 * horst_half_dx

    # JK horst block mid lower bend
    bend_horst_dz = horst_dz * maths.tan(bend_theta_di)
    cp[:, :, ni - (ni_tail + ni_bend // 2 + 1):ni - (ni_tail + ni_bend // 2 - 1), :, :, :, :] -= (horst_dz, 0.0,
                                                                                                  bend_horst_dz)
    cp[:, :, ni - (ni_tail + ni_bend // 2 - 1):, :, :, :, 2] -= 2.0 * bend_horst_dz

    faulted_grid = rqi.grid_from_cp(model,
                                    cp,
                                    crs.uuid,
                                    max_z_void = 0.01,
                                    split_pillars = True,
                                    split_tolerance = 0.01,
                                    ijk_handedness = 'right' if grid.grid_is_right_handed else 'left',
                                    known_to_be_straight = True)

    faulted_grid.write_hdf5_from_caches()
    faulted_grid.create_xml()

    # block wells against faulted grid

    log.info('faulted grid blocking of well ' + str(rqw.well_name(trajectory)))
    fbw = rqw.BlockedWell(model, grid = faulted_grid, trajectory = trajectory)
    fbw.write_hdf5()
    fbw.create_xml()
    assert fbw.cell_count == 32

    log.info('faulted grid blocking of well ' + str(rqw.well_name(traj_2)))
    fbw_2 = rqw.BlockedWell(model, grid = faulted_grid, trajectory = traj_2)
    fbw_2.write_hdf5()
    fbw_2.create_xml()
    assert fbw_2.cell_count == 26

    log.info('faulted grid blocking of well ' + str(rqw.well_name(traj_3)))
    fbw_3 = rqw.BlockedWell(model, grid = faulted_grid, trajectory = traj_3)
    fbw_3.write_hdf5()
    fbw_3.create_xml()
    assert fbw_3.cell_count == 14

    log.info('faulted grid blocking of well ' + str(rqw.well_name(traj_4)))
    fbw_4 = rqw.BlockedWell(model, grid = faulted_grid, trajectory = traj_4)
    fbw_4.write_hdf5()
    fbw_4.create_xml()
    assert fbw_4.cell_count == 16

    # create a version of the faulted grid with a k gap

    k_gap_grid = rqi.grid_from_cp(model,
                                  cp,
                                  crs.uuid,
                                  max_z_void = 0.01,
                                  split_pillars = True,
                                  split_tolerance = 0.01,
                                  ijk_handedness = 'right' if grid.grid_is_right_handed else 'left',
                                  known_to_be_straight = True)

    # convert second layer to a K gap
    k_gap_grid.nk -= 1
    k_gap_grid.extent_kji[0] = k_gap_grid.nk
    k_gap_grid.k_gaps = 1
    k_gap_grid.k_gap_after_array = np.zeros(k_gap_grid.nk - 1, dtype = bool)
    k_gap_grid.k_gap_after_array[0] = True
    k_gap_grid.k_raw_index_array = np.zeros(k_gap_grid.nk, dtype = int)
    for k in range(1, k_gap_grid.nk):
        k_gap_grid.k_raw_index_array[k] = k + 1

    # clear some attributes which may no longer be valid
    k_gap_grid.pinchout = None
    k_gap_grid.inactive = None
    k_gap_grid.grid_skin = None
    if hasattr(k_gap_grid, 'array_thickness'):
        delattr(k_gap_grid, 'array_thickness')

    k_gap_grid.write_hdf5_from_caches()
    k_gap_grid.create_xml()

    k_gap_grid_uuid = k_gap_grid.uuid

    # reload k gap grid object to ensure it is properly initialised

    k_gap_grid = None
    k_gap_grid = grr.Grid(model, uuid = k_gap_grid_uuid)

    # block wells against faulted grid with k gap

    log.info('k gap grid blocking of well ' + str(rqw.well_name(trajectory)))
    try:
        gbw = rqw.BlockedWell(model, grid = k_gap_grid, trajectory = trajectory)
        gbw.write_hdf5()
        gbw.create_xml()
        assert gbw.cell_count == 24
    except Exception:
        log.exception('failed to block well against k gap grid')

    log.info('k gap grid blocking of well ' + str(rqw.well_name(traj_2)))
    try:
        gbw_2 = rqw.BlockedWell(model, grid = k_gap_grid, trajectory = traj_2)
        gbw_2.write_hdf5()
        gbw_2.create_xml()
        assert gbw_2.cell_count == 20
    except Exception:
        log.exception('failed to block well against k gap grid')

    log.info('k gap grid blocking of well ' + str(rqw.well_name(traj_3)))
    try:
        gbw_3 = rqw.BlockedWell(model, grid = k_gap_grid, trajectory = traj_3)
        gbw_3.write_hdf5()
        gbw_3.create_xml()
        assert gbw_3.cell_count == 10
    except Exception:
        log.exception('failed to block well against k gap grid')

    log.info('k gap grid blocking of well ' + str(rqw.well_name(traj_4)))
    try:
        gbw_4 = rqw.BlockedWell(model, grid = k_gap_grid, trajectory = traj_4)
        gbw_4.write_hdf5()
        gbw_4.create_xml()
        assert gbw_4.cell_count == 10
    except Exception:
        log.exception('failed to block well against k gap grid')

    # store model

    model.store_epc()

    assert k_gap_grid.k_gaps
    assert len(k_gap_grid.k_raw_index_array) == k_gap_grid.nk
    assert len(k_gap_grid.k_gap_after_array) == k_gap_grid.nk - 1
    assert k_gap_grid.pinched_out((1, 0, 2))

    # clean up
    model.h5_release()
    os.remove(model.h5_file_name())
