# coding: utf-8
"""
This module contains the NeuberCorrection class, which is used to correct the
stress using the Neuber correction.
"""
import math
from typing import List

import matplotlib.pyplot as plt
import numpy as np

__version__ = "0.1.0"


class NeuberCorrection:
    """
    Class to correct the stress using the Neuber correction.
    It is a iterative process, so we need to correct the stress until the
    difference is less than the tolerance

    Args:
        yield_strength: The yield strength of the material
        sigma_u: The tensile strength of the material
        elastic_mod: The Young's modulus of the material
        eps_u: The strain at the tensile strength
        yield_offset: The offset of the yield strength, usually 0.002
        tolerance: The tolerance of the correction
        max_iterations: The maximum number of iterations
    """

    def __init__(
        self,
        yield_strength: float,
        sigma_u: float,
        elastic_mod: float,
        eps_u: float,
        *,
        yield_offset: float = 0.002,
        tolerance: float = 1e-6,
        max_iterations: int = 10000,
    ):
        # Validate input parameters
        if yield_strength <= 0:
            raise ValueError("yield_strength (yield strength) must be positive")
        if sigma_u <= 0:
            raise ValueError("sigma_u (tensile strength) must be positive")
        if elastic_mod <= 0:
            raise ValueError("elastic_mod (Young's modulus) must be positive")
        if eps_u <= 0:
            raise ValueError("facture_strain (strain at UTS) must be positive")
        if tolerance <= 0:
            raise ValueError("tolerance must be positive")
        if max_iterations <= 0:
            raise ValueError("max_iterations must be positive")

        # Validate that tensile strength is greater than yield strength
        if sigma_u <= yield_strength:
            raise ValueError(
                "sigma_u (tensile strength) must be greater than "
                "yield_strength (yield strength)"
            )

        self.tolerance = tolerance
        self.max_iterations = max_iterations
        self.yield_strength = yield_strength
        self.sigma_u = sigma_u
        self.elastic_mod = elastic_mod
        self.eps_u = eps_u
        self.yield_offset = yield_offset

    def _calculate_neuber_correction(self, stress: float) -> float:
        """
        Calculates the Neuber correction with Ramberg-Osgood equation.
        Uses Newton-Raphson iteration to find the corrected stress.
        """

        # Calculate Ramberg-Osgood parameter n
        elastic_ultimate = self.sigma_u / self.elastic_mod
        plastic_ultimate = self.eps_u - elastic_ultimate
        n = (math.log(plastic_ultimate) - math.log(0.002)) / math.log(
            self.sigma_u / self.yield_strength
        )

        stress_corr = stress

        for _ in range(self.max_iterations):
            # Calculate total strain (elastic + plastic)
            elastic_strain = stress_corr / self.elastic_mod
            plastic_strain = (
                self.yield_offset * (stress_corr / self.yield_strength) ** n
                if stress_corr > self.yield_strength
                else 0.0
            )
            total_strain = elastic_strain + plastic_strain

            # Calculate Neuber strain
            neuber_strain = (stress**2) / (self.elastic_mod * stress_corr)

            # Check convergence
            difference = total_strain - neuber_strain
            if abs(difference) < self.tolerance:
                break

            # Calculate derivatives for Newton-Raphson
            d_elastic = 1 / self.elastic_mod
            d_plastic = (
                self.yield_offset
                * n
                * (stress_corr / self.yield_strength) ** (n - 1)
                / self.yield_strength
                if stress_corr > self.yield_strength
                else 0.0
            )
            d_total = d_elastic + d_plastic
            d_neuber = -(stress**2) / (self.elastic_mod * stress_corr**2)

            # Newton-Raphson update
            derivative = d_total - d_neuber
            if abs(derivative) > 1e-10:
                stress_corr = stress_corr - difference / derivative
            else:
                # Fallback to simple bisection
                stress_corr = stress_corr * (0.99 if difference > 0 else 1.01)

        return stress_corr

    def correct_stress_values(self, stress_values: List[float]) -> List[float]:
        """
        Calculates the Neuber correction for the given stress.
        """
        return [self._calculate_neuber_correction(stress) for stress in stress_values]

    def correct_stress_value(self, stress: float) -> float:
        """
        Calculates the Neuber correction for the given stress.
        """
        return self._calculate_neuber_correction(stress)

    def plot_neuber_diagram(
        self,
        stress_value: float,
        show_plot: bool = True,
        save_plot: bool = False,
        plot_name: str = "Neuber_diagram",
    ):
        """
        Plots the Neuber diagram showing the Neuber hyperbolic curve and
        Ramberg-Osgood stress-strain curve for the given stress value.

        Args:
            stress_value: The elastic stress value to analyze
            show_plot: Whether to display the plot (default: True)
            save_plot: Whether to save the plot (default: False)
            plot_name: The name of the plot (default: "Neuber_diagram")
        """
        # Calculate Ramberg-Osgood parameter n
        elastic_ultimate = self.sigma_u / self.elastic_mod
        plastic_ultimate = self.eps_u - elastic_ultimate
        n = (math.log(plastic_ultimate) - math.log(0.002)) / math.log(
            self.sigma_u / self.yield_strength
        )

        # Calculate the corrected stress
        corrected_stress = self.correct_stress_value(stress_value)

        # Create stress range for plotting (avoid zero for log scale)
        stress_range = np.linspace(1, max(stress_value * 1.2, self.sigma_u), 1000)

        # Calculate Ramberg-Osgood strain for each stress
        strains_ro = []
        for stress in stress_range:
            elastic_strain = stress / self.elastic_mod
            plastic_strain = (
                self.yield_offset * (stress / self.yield_strength) ** n
                if stress > self.yield_strength
                else 0.0
            )
            strains_ro.append(elastic_strain + plastic_strain)

        # Calculate Neuber hyperbolic curve
        # Neuber's rule: σ * ε = σ_elastic^2 / E
        neuber_constant = stress_value**2 / self.elastic_mod
        strains_neuber = neuber_constant / stress_range

        # Create the plot
        fig, ax = plt.subplots(figsize=(12, 8))

        # Plot Ramberg-Osgood curve
        ax.plot(
            strains_ro,
            stress_range,
            "b-",
            linewidth=2,
            label="Ramberg-Osgood Stress-Strain Curve",
        )

        # Plot Hooke's law line (elastic line)
        hooke_strains = np.linspace(0, self.eps_u, 100)
        hooke_stresses = hooke_strains * self.elastic_mod
        ax.plot(
            hooke_strains,
            hooke_stresses,
            "g-",
            linewidth=1.5,
            alpha=0.7,
            label="Hooke's Law (Elastic Line)",
        )

        # Plot Neuber hyperbolic curve
        ax.plot(
            strains_neuber,
            stress_range,
            "r--",
            linewidth=2,
            label="Neuber Hyperbolic Curve",
        )

        # Mark the elastic stress point
        elastic_strain = stress_value / self.elastic_mod
        ax.plot(
            elastic_strain,
            stress_value,
            "go",
            markersize=10,
            label=f"Elastic Point (σ={stress_value:.0f} MPa)",
        )

        # Mark the corrected stress point
        corrected_strain = corrected_stress / self.elastic_mod
        if corrected_stress > self.yield_strength:
            plastic_strain = (
                self.yield_offset * (corrected_stress / self.yield_strength) ** n
            )
            corrected_strain += plastic_strain

        ax.plot(
            corrected_strain,
            corrected_stress,
            "mo",
            markersize=10,
            label=f"Corrected Point (σ={corrected_stress:.0f} MPa)",
        )

        # Mark yield strength line
        ax.axhline(
            y=self.yield_strength,
            color="orange",
            linestyle=":",
            label=f"Yield Strength ({self.yield_strength:.0f} MPa)",
        )

        # Mark tensile strength line
        ax.axhline(
            y=self.sigma_u,
            color="purple",
            linestyle=":",
            label=f"Tensile Strength ({self.sigma_u:.0f} MPa)",
        )

        # Add intersection point
        intersection_strain = neuber_constant / corrected_stress
        ax.plot(
            intersection_strain,
            corrected_stress,
            "ko",
            markersize=8,
            label="Intersection Point",
        )

        # Add vertical and horizontal lines at intersection point
        ax.axvline(
            x=intersection_strain,
            color="black",
            linestyle="-",
            alpha=0.5,
            linewidth=1,
        )
        ax.axhline(
            y=corrected_stress,
            color="black",
            linestyle="-",
            alpha=0.5,
            linewidth=1,
        )

        # Add value annotations
        ax.annotate(
            f"σ = {corrected_stress:.0f} MPa",
            xy=(intersection_strain, corrected_stress),
            xytext=(intersection_strain + 0.005, corrected_stress + 50),
            fontsize=10,
            ha="left",
            va="bottom",
            bbox={"boxstyle": "round,pad=0.3", "facecolor": "white", "alpha": 0.8},
            arrowprops={"arrowstyle": "->", "connectionstyle": "arc3,rad=0"},
        )

        ax.annotate(
            f"ε = {intersection_strain:.4f}",
            xy=(intersection_strain, corrected_stress),
            xytext=(intersection_strain + 0.005, corrected_stress - 100),
            fontsize=10,
            ha="left",
            va="top",
            bbox={"boxstyle": "round,pad=0.3", "facecolor": "white", "alpha": 0.8},
            arrowprops={"arrowstyle": "->", "connectionstyle": "arc3,rad=0"},
        )

        # Customize the plot
        ax.set_xlabel("Strain ε [-]", fontsize=12)
        ax.set_ylabel("Stress σ [MPa]", fontsize=12)
        ax.set_title(
            f"Neuber Diagram Analysis\n"
            f"Elastic Stress: {stress_value:.0f} MPa → "
            f"Corrected Stress: {corrected_stress:.0f} MPa",
            fontsize=14,
            fontweight="bold",
        )
        ax.grid(True, alpha=0.3)
        ax.legend(loc="upper right", fontsize=10)

        # Set reasonable axis limits
        ax.set_xlim(0, self.eps_u * 1.1)
        ax.set_ylim(0, max(stress_range) * 1.05)

        plt.tight_layout()

        if save_plot:
            plt.savefig(f"{plot_name}.png")

        if show_plot:
            plt.show()

        return fig, ax
