Calculation.py

from typing import List, Literal, Tuple

import numpy as np

from .Absolute import Absolute
from .MeasurementType import (
    MeasurementType as mt,
    DECLINATION_TYPES,
    INCLINATION_TYPES,
    MARK_TYPES,
)
from .Measurement import AverageMeasurement, Measurement, average_measurement
from .Diagnostics import Diagnostics
from .Reading import Reading


def calculate(reading: Reading, adjust_reference: bool = True) -> Reading:
    """Calculate absolutes and scale value using residual method.

    Parameters
    -------
    reading: reading to calculate absolutes from.

    Returns
    -------
    new reading object with calculated absolutes and scale_value.
    NOTE: rest of reading object is shallow copy.
    """
    # reference measurement, used to adjust absolutes
    missing_types = reading.get_missing_measurement_types()
    if len(missing_types) != 0:
        missing_types = ", ".join(t.value for t in missing_types)
        raise ValueError(f"Missing {missing_types} measurements in input reading")
    reference = adjust_reference and reading[mt.WEST_DOWN][0] or None
    # calculate inclination
    inclination, f, i_mean = calculate_I(
        hemisphere=reading.hemisphere, measurements=reading.measurements
    )
    corrected_f = f + reading.pier_correction

    # calculate absolutes,
    (absoluteH, absoluteZ) = calculate_HZ_absolutes(
        corrected_f=corrected_f,
        inclination=inclination,
        mean=i_mean,
        reference=reference,
    )

    horizontal_component = calculate_horizontal_component(
        inclination=inclination, corrected_f=corrected_f
    )

    vertical_component = calculate_vertical_component(
        inclination=inclination, corrected_f=corrected_f
    )

    absoluteD, meridian = calculate_D_absolute(
        azimuth=reading.azimuth,
        h_baseline=absoluteH.baseline,
        measurements=reading.measurements,
        reference=reference,
        shift=reading.absolutes[0].shift,
    )

    d_computed = calculate_D_computed(
        measurements=reading.measurements,
        h_baseline=absoluteH.baseline,
    )

    r_computed = calculate_R_computed(
        measurements=reading.measurements,
        h_baseline=absoluteH.baseline,
    )

    magnetic_south_meridian = calculate_magnetic_south_meridian(
        measurements=reading.measurements,
    )

    # calculate scale
    scale_value = None
    scale_measurements = reading[mt.NORTH_DOWN_SCALE]
    if scale_measurements:
        scale_value = calculate_scale_value(
            corrected_f=corrected_f,
            inclination=inclination,
            measurements=scale_measurements,
        )

    mean_mark = calculate_mean_mark(
        measurements=reading.measurements,
    )

    magnetic_azimuth_of_mark = calculate_magnetic_azimuth_of_mark(
        mean_mark=mean_mark, magnetic_south_meridian=magnetic_south_meridian
    )

    ordinate_f = f

    ordinate_d = calculate_ordinate_d(
        absoluteD_absolute=absoluteD.absolute,
        absoluteD_baseline=absoluteD.baseline,
    )

    ordinate_h = calculate_ordinate_h(
        absoluteH_absolute=absoluteH.absolute,
        absoluteH_baseline=absoluteH.baseline,
        ordinate_d=ordinate_d,
    )

    ordinate_z = calculate_ordinate_z(
        absoluteZ_absolute=absoluteZ.absolute,
        absoluteZ_baseline=absoluteZ.baseline,
    )

    # populate diagnostics object with averaged measurements
    diagnostics = Diagnostics(
        corrected_f=corrected_f,
        d_computed=d_computed,
        horizontal_component=horizontal_component,
        inclination=inclination,
        magnetic_azimuth_of_mark=magnetic_azimuth_of_mark,
        magnetic_south_meridian=magnetic_south_meridian,
        mean_mark=mean_mark,
        ordinate_d=ordinate_d,
        ordinate_f=ordinate_f,
        ordinate_h=ordinate_h,
        ordinate_z=ordinate_z,
        r_computed=r_computed,
        vertical_component=vertical_component,
    )

    # create new reading object
    calculated = Reading(
        absolutes=[absoluteD, absoluteH, absoluteZ],
        diagnostics=diagnostics,
        scale_value=scale_value,
        # copy other attributes
        **reading.model_dump(exclude={"absolutes", "diagnostics", "scale_value"}),
    )

    return calculated


def calculate_D_absolute(
    measurements: List[Measurement],
    azimuth: float,
    h_baseline: float,
    reference: Measurement,
    shift: float = 0,
) -> Tuple[Absolute, Diagnostics]:
    """Calculate D absolute.

    Parameters
    ----------
    measurements: list with at least declination and mark measurements.
    azimuth: azimuth to mark in decimal degrees.
    h_baseline: calculated H baseline value.
    reference: reference measurement used to adjust absolute.
    shift: usually 0 or +/-180 for Declination

    Returns
    -------
    D Absolute
    """
    mean = average_measurement(measurements, DECLINATION_TYPES)
    if reference:
        starttime = reference.time
        endtime = reference.time
    else:
        starttime = mean.time
        endtime = mean.endtime
    reference = reference or mean
    # average mark
    average_mark = average_measurement(measurements, MARK_TYPES)
    # adjust based on which is larger
    mark_up = average_measurement(measurements, [mt.FIRST_MARK_UP]).angle
    mark_down = average_measurement(measurements, [mt.FIRST_MARK_DOWN]).angle
    if mark_up < mark_down:
        average_mark.angle += 90
    else:
        average_mark.angle -= 90
    # declination measurements
    declination_measurements = [
        average_measurement(measurements, [t]) for t in DECLINATION_TYPES
    ]
    # average declination meridian
    meridian = np.average(
        [
            m.angle
            + np.degrees(
                m.measurement_type.meridian
                * (np.arcsin(m.residual / np.sqrt((m.h + h_baseline) ** 2 + m.e**2)))
            )
            - np.degrees(np.arctan(m.e / (m.h + h_baseline)))
            for m in declination_measurements
        ]
    )

    # add subtract average mark angle from average meridian angle and add
    # azimuth to get the declination baseline
    d_b = (meridian - average_mark.angle) + azimuth

    # given operational limitations, there is always an 180-degree ambiguity
    # in our calculated d_b; therefore, we have to assume the magnetic field
    # points north(ish), and force d_b to fall between -90 and +90 degrees
    while d_b < -90:
        d_b += 180
    while d_b > 90:
        d_b -= 180

    # the preceeding lead to a wrong answer if the magnetic field actually
    # points south(ish);
    d_b += shift

    # (for example, there are places in Antarctica where the magnetic field
    #  actually points south-ish)

    # calculate absolute
    d_abs = d_b + np.degrees(np.arctan(reference.e / (reference.h + h_baseline)))

    return (
        Absolute(
            element="D",
            absolute=d_abs,
            baseline=d_b,
            shift=shift,
            starttime=starttime,
            endtime=endtime,
        ),
        meridian,
    )


def calculate_HZ_absolutes(
    inclination: float,
    corrected_f: float,
    mean: AverageMeasurement,
    reference: Measurement,
) -> Tuple[Absolute, Absolute, float, float]:
    """Calculate H and Z absolutes

    Parameters
    ----------
    inclination: calculated inclination.
    corrected_f: calculated f with pier correction.
    mean: mean of inclination ordinates.
    reference:  reference measurement used to adjust absolutes.

    Returns
    -------
    Tuple
        - H Absolute
        - Z Absolute
    """
    inclination_radians = np.radians(inclination)
    h_abs = corrected_f * np.cos(inclination_radians)
    z_abs = corrected_f * np.sin(inclination_radians)
    h_b = np.sqrt(h_abs**2 - mean.e**2) - mean.h
    z_b = z_abs - mean.z
    # adjust absolutes to reference measurement
    if reference:
        h_abs = np.sqrt((h_b + reference.h) ** 2 + (reference.e) ** 2)
        z_abs = z_b + reference.z
        starttime = reference.time
        endtime = reference.time
    else:
        starttime = mean.time
        endtime = mean.endtime
    # return absolutes
    return (
        Absolute(
            element="H",
            baseline=h_b,
            absolute=h_abs,
            starttime=starttime,
            endtime=endtime,
        ),
        Absolute(
            element="Z",
            baseline=z_b,
            absolute=z_abs,
            starttime=starttime,
            endtime=endtime,
        ),
    )


def calculate_I(
    measurements: List[Measurement], hemisphere: Literal[-1, 1] = 1
) -> Tuple[float, float, Measurement]:
    """Calculate inclination and f from measurements.

    Parameters
    ----------
    measurements: list with at least inclination measurements.
    hemisphere: +1 for northern hemisphere (default), -1 for southern hemisphere.

    Returns
    -------
    Tuple
        - inclination angle in decimal degrees
        - uncorrected calculated f
        - mean of inclination measurements
    """
    # mean across all inclination measurements
    mean = average_measurement(measurements, INCLINATION_TYPES)
    # mean within each type
    inclination_measurements = [
        average_measurement(measurements, [t]) for t in INCLINATION_TYPES
    ]
    # get initial inclination angle, assumed to be the southdown angle
    inclination = average_measurement(measurements, [mt.SOUTH_DOWN]).angle
    if inclination >= 90:
        inclination -= 180
    # loop until inclination converges
    last_inclination = inclination + 1
    while abs(last_inclination - inclination) > 0.0001:
        # set temporary inclination variable to hold previous step's inclination
        last_inclination = inclination
        # Update measurement f based on inclination
        inclination_radians = np.radians(inclination)
        for m in inclination_measurements:
            m.f = (
                mean.f
                + (m.h - mean.h) * np.cos(inclination_radians)
                + (m.z - mean.z) * np.sin(inclination_radians)
                + ((m.e) ** 2 - (mean.e) ** 2) / (2 * mean.f)
            )
        # calculate average inclination
        inclination = np.average(
            [
                (
                    m.measurement_type.shift
                    + m.measurement_type.meridian
                    * (
                        +m.angle
                        + m.measurement_type.direction
                        * (hemisphere * np.degrees(np.arcsin(m.residual / m.f)))
                    )
                )
                for m in inclination_measurements
            ]
        )
    # calculate uncorrected f
    f = np.average([m.f for m in inclination_measurements])
    return inclination, f, mean


def calculate_scale_value(
    measurements: List[Measurement], inclination: float, corrected_f: float
) -> float:
    """Calculate scale value.

    Parameters
    ----------
    measurements: measurements to be used for scale calculation.
        should have type NORTH_DOWN_SCALE.
    inclination: calculated inclination.
    corrected_f: calculated f with pier correction.

    Returns
    -------
    Calculated scale value.
    """
    inclination_radians = np.radians(inclination)
    m1, m2 = measurements[0], measurements[-1]
    field_change = np.degrees(
        (
            -np.sin(inclination_radians) * (m2.h - m1.h)
            + np.cos(inclination_radians) * (m2.z - m1.z)
        )
        / corrected_f
    ) + (m2.angle - m1.angle)
    residual_change = m2.residual - m1.residual
    scale_value = np.radians(corrected_f * field_change / np.abs(residual_change))
    return scale_value


def calculate_magnetic_south_meridian(measurements: List[Measurement]) -> float:
    """Calculate magnetic_south_meridian.

    Parameters
    ----------
    measurements: measurements, should have type WEST_DOWN, EAST_DOWN, WEST_UP, EAST_UP.

    Returns
    -------
    Calculated magnetic_south_meridian.
    """
    # declination measurements
    declination_measurements = [
        average_measurement(measurements, [t]) for t in DECLINATION_TYPES
    ]

    magnetic_south_meridian = np.average([m.angle for m in declination_measurements])
    return magnetic_south_meridian


def calculate_D_computed(
    measurements: List[Measurement],
    h_baseline: float,
) -> float:
    """Calculate d_computed.

    Parameters
    ----------
    measurements: measurements,
    h_baseline

    Returns
    -------
    Calculated d_computed.
    """
    # declination measurements
    declination_measurements = [
        average_measurement(measurements, [t]) for t in DECLINATION_TYPES
    ]

    d_computed = np.average(
        [
            -1 * np.degrees(np.arctan(m.e / (m.h + h_baseline)))
            for m in declination_measurements
        ]
    )

    return d_computed


def calculate_R_computed(
    measurements: List[Measurement],
    h_baseline: float,
) -> float:
    """Calculate r_computed.

    Parameters
    ----------
    measurements: measurements,
    h_baseline

    Returns
    -------
    Calculated r_computed.
    """
    # declination measurements
    declination_measurements = [
        average_measurement(measurements, [t]) for t in DECLINATION_TYPES
    ]

    r_computed = np.average(
        [
            np.degrees(
                m.measurement_type.meridian
                * (np.arcsin(m.residual / np.sqrt((m.h + h_baseline) ** 2 + m.e**2)))
            )
            for m in declination_measurements
        ]
    )

    return r_computed


def calculate_mean_mark(
    measurements: List[Measurement],
) -> float:
    """Calculate mean_mark.

    Parameters
    ----------
    measurements: measurements,


    Returns
    -------
    Calculated mean_mark angle.
    """

    # average mark
    average_mark = average_measurement(measurements, MARK_TYPES)

    return average_mark.angle


def calculate_magnetic_azimuth_of_mark(
    mean_mark: float,
    magnetic_south_meridian: float,
) -> float:
    """Calculate magnetic_azimuth_of_mark.

    Parameters
    ----------
    mean_mark: float,
    magnetic_south_meridian:float,


    Returns
    -------
    Calculated magnetic_azimuth_of_mark
    """

    magnetic_azimuth_of_mark = mean_mark - magnetic_south_meridian

    return magnetic_azimuth_of_mark


def calculate_ordinate_d(
    absoluteD_absolute: float,
    absoluteD_baseline: float,
) -> float:
    """Calculate ordinate_d.

    Parameters
    ----------
    absoluteD_absolute: float,
    absoluteD_baseline:float,


    Returns
    -------
    Calculated ordinate_d
    """

    ordinate_d = absoluteD_absolute - absoluteD_baseline

    return ordinate_d


def calculate_ordinate_z(
    absoluteZ_absolute: float,
    absoluteZ_baseline: float,
) -> float:
    """Calculate ordinate_z.

    Parameters
    ----------
    absoluteZ_absolute: float,
    absoluteZ_baseline:float,


    Returns
    -------
    Calculated ordinate_z
    """

    ordinate_z = absoluteZ_absolute - absoluteZ_baseline

    return ordinate_z


def calculate_ordinate_h(
    absoluteH_absolute: float,
    absoluteH_baseline: float,
    ordinate_d: float,
) -> float:
    """Calculate ordinate_z.

    Parameters
    ----------
    absoluteH_absolute: float,
    absoluteH_baseline:float,
    ordinate_d:float,

    Returns
    -------
    Calculated ordinate_h
    """
    ordinate_e = absoluteH_absolute * np.sin(np.radians(ordinate_d))

    ordinate_h = np.sqrt(absoluteH_absolute**2 - ordinate_e**2) - absoluteH_baseline

    return ordinate_h


def calculate_horizontal_component(inclination: float, corrected_f: float) -> float:
    """Calculate horizontal_component.

    Parameters
    ----------
    inclination: calculated inclination.
    corrected_f: calculated f with pier correction.

    Returns
    -------
    Calculated horizontal_component
    """

    inclination_radians = np.radians(inclination)
    horizontal_component = corrected_f * np.cos(inclination_radians)

    return horizontal_component


def calculate_vertical_component(inclination: float, corrected_f: float) -> float:
    """Calculate vertical_component.

    Parameters
    ----------
    inclination: calculated inclination.
    corrected_f: calculated f with pier correction.

    Returns
    -------
    Calculated vertical_component
    """

    inclination_radians = np.radians(inclination)
    vertical_component = corrected_f * np.sin(inclination_radians)

    return vertical_component