#!/usr/bin/env python3.10

"""Estimate Wilson's position with Newton's method."""

from itertools import combinations
from pathlib import Path
import shutil

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np


def evaluate_system(state_vector, first_ship, second_ship, third_ship, ship_data):
    """Evaluate the nonlinear system associated with three selected ships."""
    light_speed = 3.0e5
    state_vector = np.asarray(state_vector, dtype=float).reshape(3)
    residual_vector = np.zeros(3, dtype=float)

    residual_vector[0] = (
        (state_vector[0] - ship_data[first_ship, 0]) ** 2
        + (state_vector[1] - ship_data[first_ship, 1]) ** 2
        - (light_speed * (ship_data[first_ship, 2] - state_vector[2])) ** 2
    )
    residual_vector[1] = (
        (state_vector[0] - ship_data[second_ship, 0]) ** 2
        + (state_vector[1] - ship_data[second_ship, 1]) ** 2
        - (light_speed * (ship_data[second_ship, 2] - state_vector[2])) ** 2
    )
    residual_vector[2] = (
        (state_vector[0] - ship_data[third_ship, 0]) ** 2
        + (state_vector[1] - ship_data[third_ship, 1]) ** 2
        - (light_speed * (ship_data[third_ship, 2] - state_vector[2])) ** 2
    )
    return residual_vector


def evaluate_jacobian(state_vector, first_ship, second_ship, third_ship, ship_data):
    """Evaluate the Jacobian matrix of the nonlinear system."""
    light_speed = 3.0e5
    state_vector = np.asarray(state_vector, dtype=float).reshape(3)
    jacobian = np.zeros((3, 3), dtype=float)

    jacobian[0, 0] = state_vector[0] - ship_data[first_ship, 0]
    jacobian[0, 1] = state_vector[1] - ship_data[first_ship, 1]
    jacobian[0, 2] = light_speed**2 * (ship_data[first_ship, 2] - state_vector[2])

    jacobian[1, 0] = state_vector[0] - ship_data[second_ship, 0]
    jacobian[1, 1] = state_vector[1] - ship_data[second_ship, 1]
    jacobian[1, 2] = light_speed**2 * (ship_data[second_ship, 2] - state_vector[2])

    jacobian[2, 0] = state_vector[0] - ship_data[third_ship, 0]
    jacobian[2, 1] = state_vector[1] - ship_data[third_ship, 1]
    jacobian[2, 2] = light_speed**2 * (ship_data[third_ship, 2] - state_vector[2])

    return 2.0 * jacobian


def solve_newton_system(initial_guess, ship_indices, ship_data, tolerance, max_iterations):
    """Solve one three-equation nonlinear system with Newton's method."""
    current_solution = initial_guess.copy()

    for _ in range(max_iterations):
        jacobian_matrix = evaluate_jacobian(current_solution, *ship_indices, ship_data)
        residual_vector = evaluate_system(current_solution, *ship_indices, ship_data)
        next_solution = current_solution - np.linalg.solve(jacobian_matrix, residual_vector)

        denominator = np.linalg.norm(next_solution, ord=np.inf)
        if denominator == 0.0:
            return next_solution

        relative_error = np.linalg.norm(next_solution - current_solution, ord=np.inf) / denominator
        current_solution = next_solution
        if relative_error <= tolerance:
            return current_solution

    return current_solution


def build_ship_data():
    """Build the ship-position and timing dataset used in the exercise."""
    x_coordinates = np.array([47.0, 90.0, 140.0, 190.0], dtype=float)
    y_coordinates = np.array([102.0, 42.0, 101.0, 48.0], dtype=float)
    time_deltas = 1e-6 * np.array([440.15, 277.14, 236.67, 271.40], dtype=float)
    ship_data = np.column_stack((x_coordinates, y_coordinates, time_deltas))
    return x_coordinates, y_coordinates, ship_data


def configure_plot_style():
    """Configure a consistent LaTeX-like plotting style."""
    mpl.rcParams.update(
        {
            "text.usetex": shutil.which("latex") is not None,
            "font.family": "serif",
            "font.serif": ["Computer Modern Roman", "DejaVu Serif"],
            "mathtext.fontset": "cm",
            "axes.titlesize": 15,
            "axes.labelsize": 13,
            "xtick.labelsize": 11,
            "ytick.labelsize": 11,
            "legend.fontsize": 11,
        }
    )


def main():
    """Run the full location-recovery experiment."""
    configure_plot_style()
    output_directory = Path(__file__).resolve().parent.parent / "Figuras"
    output_directory.mkdir(exist_ok=True)

    ship_count = 4
    x_coordinates, y_coordinates, ship_data = build_ship_data()
    tolerance = 1.0e-7
    max_iterations = 200
    initial_guess = np.array([150.0, 80.0, 0.0], dtype=float)

    candidate_solutions = []
    for ship_indices in combinations(range(ship_count), 3):
        candidate_solutions.append(
            solve_newton_system(initial_guess, ship_indices, ship_data, tolerance, max_iterations)
        )

    candidate_positions = np.array(candidate_solutions, dtype=float).T
    wilson_position = np.array(
        [np.mean(candidate_positions[0, :]), np.mean(candidate_positions[1, :])], dtype=float
    )

    ship_positions = np.vstack((x_coordinates, y_coordinates))
    distances_to_ships = np.linalg.norm(ship_positions.T - wilson_position, axis=1)
    nearest_ship = int(np.argmin(distances_to_ships) + 1)

    plt.figure(figsize=(7.0, 5.0))
    plt.scatter(x_coordinates, y_coordinates, marker="s", label="Ships")
    plt.scatter(wilson_position[0], wilson_position[1], marker="*", s=160, label="Wilson")
    for ship_index in range(ship_count):
        plt.text(x_coordinates[ship_index] + 2.0, y_coordinates[ship_index] + 1.5, f"N{ship_index + 1}")
    plt.xlabel("x (km)")
    plt.ylabel("y (km)")
    plt.grid(True, linewidth=0.4, alpha=0.5)
    plt.legend()
    plt.tight_layout()
    plt.savefig(output_directory / "ship_location.pdf", bbox_inches="tight")
    plt.close()

    print("Wilson position =", wilson_position)
    print("Nearest ship =", nearest_ship)


if __name__ == "__main__":
    main()