"""Thick-walled pressure vessel solvers using Lame equations.

Two configurations:
- Thick-walled cylinder under internal pressure
- Thick-walled sphere under internal pressure

These are exact analytical solutions for linear elastic, isotropic materials
under axisymmetric loading. The maximum stress occurs at the inner surface.

Reference: Boresi, A.P. & Schmidt, R.J., "Advanced Mechanics of Materials"
"""

import math
from typing import Any

from src.data.schema import SolutionResult
from src.data.solvers.base import AnalyticalSolver


class ThickCylinder(AnalyticalSolver):
    """Thick-walled cylinder under internal pressure p (Lame equations).

    Hoop stress at inner surface (maximum):
        sigma_h = p * (r_o^2 + r_i^2) / (r_o^2 - r_i^2)

    Radial stress at inner surface:
        sigma_r = -p  (compressive)

    Radial displacement at inner surface:
        u_r = (p * r_i / E) * [(1-nu)(r_i^2) + (1+nu)(r_o^2)] / (r_o^2 - r_i^2)

    Von Mises equivalent stress at inner surface (plane stress, open ends):
        sigma_vm = sqrt(sigma_h^2 + sigma_r^2 - sigma_h*sigma_r)
    """

    @property
    def config_id(self) -> str:
        return "vessel_cylinder"

    @property
    def problem_family(self) -> str:
        return "vessel"

    def solve(self, params: dict[str, Any]) -> SolutionResult:
        r_i = params["inner_radius"]
        r_o = params["outer_radius"]
        E = params["elastic_modulus"]
        nu = params["poisson_ratio"]
        sigma_y = params["yield_strength"]
        p = params["internal_pressure"]

        ri2 = r_i**2
        ro2 = r_o**2
        denom = ro2 - ri2

        # Stresses at inner surface (r = r_i)
        sigma_hoop = p * (ro2 + ri2) / denom
        sigma_radial = -p  # compressive

        # Von Mises equivalent stress (plane stress, open-ended cylinder)
        sigma_vm = math.sqrt(
            sigma_hoop**2 + sigma_radial**2 - sigma_hoop * sigma_radial
        )

        # Radial displacement at inner surface
        u_r = (p * r_i / E) * ((1.0 - nu) * ri2 + (1.0 + nu) * ro2) / denom

        return SolutionResult.from_stress(sigma_vm, abs(u_r), sigma_y)


class ThickSphere(AnalyticalSolver):
    """Thick-walled sphere under internal pressure p.

    Hoop stress at inner surface (maximum, equal in both tangential directions):
        sigma_h = p * r_i^3 * (r_o^3 + 2*r_i^3) / (2 * r_i^3 * (r_o^3 - r_i^3))

    Simplified:
        sigma_h = p * (r_o^3 + 2*r_i^3) / (2 * (r_o^3 - r_i^3))

    Radial stress at inner surface:
        sigma_r = -p

    Von Mises at inner surface (biaxial hoop stress):
        sigma_vm = sqrt(sigma_h^2 + sigma_h^2 - sigma_h*sigma_h + 3*0)
        For sigma_1 = sigma_2 = sigma_h, sigma_3 = sigma_r:
        sigma_vm = sqrt(sigma_h^2 - sigma_h*sigma_r + sigma_r^2)

    Radial displacement at inner surface:
        u_r = (p * r_i) / (2*E) * [(1-2nu)*r_i^3 + (1+nu)*r_o^3] / (r_o^3 - r_i^3)
    """

    @property
    def config_id(self) -> str:
        return "vessel_sphere"

    @property
    def problem_family(self) -> str:
        return "vessel"

    def solve(self, params: dict[str, Any]) -> SolutionResult:
        r_i = params["inner_radius"]
        r_o = params["outer_radius"]
        E = params["elastic_modulus"]
        nu = params["poisson_ratio"]
        sigma_y = params["yield_strength"]
        p = params["internal_pressure"]

        ri3 = r_i**3
        ro3 = r_o**3
        denom = ro3 - ri3

        # Stresses at inner surface
        sigma_hoop = p * (ro3 + 2.0 * ri3) / (2.0 * denom)
        sigma_radial = -p

        # Von Mises: two equal hoop stresses + radial
        # sigma_vm = sqrt(s1^2 + s2^2 + s3^2 - s1*s2 - s2*s3 - s1*s3)
        # with s1 = s2 = sigma_hoop, s3 = sigma_radial
        sigma_vm = math.sqrt(
            sigma_hoop**2 - sigma_hoop * sigma_radial + sigma_radial**2
        )

        # Radial displacement at inner surface
        u_r = (p * r_i / (2.0 * E)) * (
            (1.0 - 2.0 * nu) * ri3 + (1.0 + nu) * ro3
        ) / denom

        return SolutionResult.from_stress(sigma_vm, abs(u_r), sigma_y)


VESSEL_SOLVERS: dict[str, type[AnalyticalSolver]] = {
    "vessel_cylinder": ThickCylinder,
    "vessel_sphere": ThickSphere,
}