"""
analytical_cases.py
Analytical solutions for classical mechanics problems

Provides known solutions for validation:
- Cantilever beam under point load
- Simply supported beam
- Axial tension bar
- Pressure vessel (thin-walled cylinder)
- Torsion of circular shaft

These serve as ground truth for testing neural predictions.
"""

import numpy as np
from dataclasses import dataclass
from typing import Tuple


@dataclass
class BeamProperties:
    """Material and geometric properties for beam"""
    length: float  # m
    width: float  # m
    height: float  # m
    E: float  # Young's modulus (Pa)
    nu: float  # Poisson's ratio
    rho: float  # Density (kg/m³)

    @property
    def I(self) -> float:
        """Second moment of area for rectangular section"""
        return (self.width * self.height**3) / 12

    @property
    def A(self) -> float:
        """Cross-sectional area"""
        return self.width * self.height

    @property
    def G(self) -> float:
        """Shear modulus"""
        return self.E / (2 * (1 + self.nu))


def cantilever_beam_point_load(force: float, props: BeamProperties) -> dict:
    """
    Cantilever beam with point load at free end

    Analytical solution:
        δ_max = FL³/3EI  (at free end)
        σ_max = FL/Z     (at fixed end)

    Args:
        force: Applied force at tip (N)
        props: Beam properties

    Returns:
        dict with displacement, stress, reactions
    """
    L = props.length
    E = props.E
    I = props.I
    c = props.height / 2  # Distance to neutral axis

    # Maximum displacement at free end
    delta_max = (force * L**3) / (3 * E * I)

    # Maximum stress at fixed end (bending)
    M_max = force * L  # Maximum moment at fixed end
    Z = I / c  # Section modulus
    sigma_max = M_max / Z

    # Deflection curve: y(x) = (F/(6EI)) * x² * (3L - x)
    def deflection_at(x):
        """Deflection at position x along beam"""
        if x < 0 or x > L:
            return 0.0
        return (force / (6 * E * I)) * x**2 * (3 * L - x)

    # Bending moment: M(x) = F * (L - x)
    def moment_at(x):
        """Bending moment at position x"""
        if x < 0 or x > L:
            return 0.0
        return force * (L - x)

    # Stress: σ(x, y) = M(x) * y / I
    def stress_at(x, y):
        """Bending stress at position x and distance y from neutral axis"""
        M = moment_at(x)
        return (M * y) / I

    return {
        'type': 'cantilever_point_load',
        'delta_max': delta_max,
        'sigma_max': sigma_max,
        'deflection_function': deflection_at,
        'moment_function': moment_at,
        'stress_function': stress_at,
        'load': force,
        'properties': props
    }


def simply_supported_beam_point_load(force: float, props: BeamProperties) -> dict:
    """
    Simply supported beam with point load at center

    Analytical solution:
        δ_max = FL³/48EI  (at center)
        σ_max = FL/4Z     (at center)

    Args:
        force: Applied force at center (N)
        props: Beam properties

    Returns:
        dict with displacement, stress, reactions
    """
    L = props.length
    E = props.E
    I = props.I
    c = props.height / 2

    # Maximum displacement at center
    delta_max = (force * L**3) / (48 * E * I)

    # Maximum stress at center
    M_max = force * L / 4  # Maximum moment at center
    Z = I / c
    sigma_max = M_max / Z

    # Deflection curve (for 0 < x < L/2):
    # y(x) = (F/(48EI)) * x * (3L² - 4x²)
    def deflection_at(x):
        """Deflection at position x along beam"""
        if x < 0 or x > L:
            return 0.0
        if x <= L/2:
            return (force / (48 * E * I)) * x * (3 * L**2 - 4 * x**2)
        else:
            # Symmetric about center
            return deflection_at(L - x)

    # Bending moment: M(x) = (F/2) * x  for x < L/2
    def moment_at(x):
        """Bending moment at position x"""
        if x < 0 or x > L:
            return 0.0
        if x <= L/2:
            return (force / 2) * x
        else:
            return (force / 2) * (L - x)

    def stress_at(x, y):
        """Bending stress at position x and distance y from neutral axis"""
        M = moment_at(x)
        return (M * y) / I

    return {
        'type': 'simply_supported_point_load',
        'delta_max': delta_max,
        'sigma_max': sigma_max,
        'deflection_function': deflection_at,
        'moment_function': moment_at,
        'stress_function': stress_at,
        'load': force,
        'properties': props,
        'reactions': force / 2  # Each support
    }


def axial_tension_bar(force: float, props: BeamProperties) -> dict:
    """
    Bar under axial tension

    Analytical solution:
        δ = FL/EA  (total elongation)
        σ = F/A    (uniform stress)
        ε = σ/E    (uniform strain)

    Args:
        force: Axial force (N)
        props: Bar properties

    Returns:
        dict with displacement, stress, strain
    """
    L = props.length
    E = props.E
    A = props.A

    # Total elongation
    delta = (force * L) / (E * A)

    # Uniform stress
    sigma = force / A

    # Uniform strain
    epsilon = sigma / E

    # Displacement is linear along length
    def displacement_at(x):
        """Axial displacement at position x"""
        if x < 0 or x > L:
            return 0.0
        return (force * x) / (E * A)

    return {
        'type': 'axial_tension',
        'delta_total': delta,
        'sigma': sigma,
        'epsilon': epsilon,
        'displacement_function': displacement_at,
        'load': force,
        'properties': props
    }


def thin_wall_pressure_vessel(pressure: float, radius: float, thickness: float,
                               length: float, E: float, nu: float) -> dict:
    """
    Thin-walled cylindrical pressure vessel

    Analytical solution:
        σ_hoop = pr/t      (circumferential stress)
        σ_axial = pr/2t    (longitudinal stress)
        ε_hoop = (1/E)(σ_h - ν*σ_a)
        ε_axial = (1/E)(σ_a - ν*σ_h)

    Args:
        pressure: Internal pressure (Pa)
        radius: Mean radius (m)
        thickness: Wall thickness (m)
        length: Cylinder length (m)
        E: Young's modulus (Pa)
        nu: Poisson's ratio

    Returns:
        dict with stresses and strains
    """
    # Stresses
    sigma_hoop = (pressure * radius) / thickness
    sigma_axial = (pressure * radius) / (2 * thickness)

    # Strains
    epsilon_hoop = (1/E) * (sigma_hoop - nu * sigma_axial)
    epsilon_axial = (1/E) * (sigma_axial - nu * sigma_hoop)

    # Radial expansion
    delta_r = epsilon_hoop * radius

    return {
        'type': 'pressure_vessel',
        'sigma_hoop': sigma_hoop,
        'sigma_axial': sigma_axial,
        'epsilon_hoop': epsilon_hoop,
        'epsilon_axial': epsilon_axial,
        'radial_expansion': delta_r,
        'pressure': pressure,
        'radius': radius,
        'thickness': thickness
    }


def torsion_circular_shaft(torque: float, radius: float, length: float, G: float) -> dict:
    """
    Circular shaft under torsion

    Analytical solution:
        θ = TL/GJ          (angle of twist)
        τ_max = Tr/J       (maximum shear stress at surface)
        γ_max = τ_max/G    (maximum shear strain)

    Args:
        torque: Applied torque (N·m)
        radius: Shaft radius (m)
        length: Shaft length (m)
        G: Shear modulus (Pa)

    Returns:
        dict with twist angle, stress, strain
    """
    # Polar moment of inertia
    J = (np.pi * radius**4) / 2

    # Angle of twist
    theta = (torque * length) / (G * J)

    # Maximum shear stress (at surface)
    tau_max = (torque * radius) / J

    # Maximum shear strain
    gamma_max = tau_max / G

    # Shear stress at radius r
    def shear_stress_at(r):
        """Shear stress at radial distance r from center"""
        if r < 0 or r > radius:
            return 0.0
        return (torque * r) / J

    return {
        'type': 'torsion',
        'theta': theta,
        'tau_max': tau_max,
        'gamma_max': gamma_max,
        'shear_stress_function': shear_stress_at,
        'torque': torque,
        'radius': radius,
        'length': length
    }


# Standard test cases with typical values
def get_standard_cantilever() -> Tuple[float, BeamProperties]:
    """Standard cantilever test case"""
    props = BeamProperties(
        length=1.0,      # 1 m
        width=0.05,      # 50 mm
        height=0.1,      # 100 mm
        E=210e9,         # Steel: 210 GPa
        nu=0.3,
        rho=7850         # kg/m³
    )
    force = 1000.0  # 1 kN
    return force, props


def get_standard_simply_supported() -> Tuple[float, BeamProperties]:
    """Standard simply supported beam test case"""
    props = BeamProperties(
        length=2.0,      # 2 m
        width=0.05,      # 50 mm
        height=0.1,      # 100 mm
        E=210e9,         # Steel
        nu=0.3,
        rho=7850
    )
    force = 5000.0  # 5 kN
    return force, props


def get_standard_tension_bar() -> Tuple[float, BeamProperties]:
    """Standard tension bar test case"""
    props = BeamProperties(
        length=1.0,      # 1 m
        width=0.02,      # 20 mm
        height=0.02,     # 20 mm (square bar)
        E=210e9,         # Steel
        nu=0.3,
        rho=7850
    )
    force = 10000.0  # 10 kN
    return force, props


# Example usage and validation
if __name__ == "__main__":
    print("Analytical Test Cases\n")
    print("="*60)

    # Test 1: Cantilever beam
    print("\n1. Cantilever Beam (Point Load at Tip)")
    print("-"*60)
    force, props = get_standard_cantilever()
    result = cantilever_beam_point_load(force, props)

    print(f"Load: {force} N")
    print(f"Length: {props.length} m")
    print(f"E: {props.E/1e9:.0f} GPa")
    print(f"I: {props.I*1e12:.3f} mm⁴")
    print(f"\nResults:")
    print(f"  Max displacement: {result['delta_max']*1000:.3f} mm")
    print(f"  Max stress: {result['sigma_max']/1e6:.1f} MPa")

    # Verify deflection at intermediate points
    print(f"\nDeflection profile:")
    for x in [0.0, 0.25, 0.5, 0.75, 1.0]:
        x_m = x * props.length
        delta = result['deflection_function'](x_m)
        print(f"  x = {x:.2f}L: δ = {delta*1000:.3f} mm")

    # Test 2: Simply supported beam
    print("\n2. Simply Supported Beam (Point Load at Center)")
    print("-"*60)
    force, props = get_standard_simply_supported()
    result = simply_supported_beam_point_load(force, props)

    print(f"Load: {force} N")
    print(f"Length: {props.length} m")
    print(f"\nResults:")
    print(f"  Max displacement: {result['delta_max']*1000:.3f} mm")
    print(f"  Max stress: {result['sigma_max']/1e6:.1f} MPa")
    print(f"  Reactions: {result['reactions']} N each")

    # Test 3: Axial tension
    print("\n3. Axial Tension Bar")
    print("-"*60)
    force, props = get_standard_tension_bar()
    result = axial_tension_bar(force, props)

    print(f"Load: {force} N")
    print(f"Length: {props.length} m")
    print(f"Area: {props.A*1e6:.0f} mm²")
    print(f"\nResults:")
    print(f"  Total elongation: {result['delta_total']*1e6:.3f} μm")
    print(f"  Stress: {result['sigma']/1e6:.1f} MPa")
    print(f"  Strain: {result['epsilon']*1e6:.1f} με")

    # Test 4: Pressure vessel
    print("\n4. Thin-Walled Pressure Vessel")
    print("-"*60)
    pressure = 10e6  # 10 MPa
    radius = 0.5     # 500 mm
    thickness = 0.01  # 10 mm
    result = thin_wall_pressure_vessel(pressure, radius, thickness, 2.0, 210e9, 0.3)

    print(f"Pressure: {pressure/1e6:.1f} MPa")
    print(f"Radius: {radius*1000:.0f} mm")
    print(f"Thickness: {thickness*1000:.0f} mm")
    print(f"\nResults:")
    print(f"  Hoop stress: {result['sigma_hoop']/1e6:.1f} MPa")
    print(f"  Axial stress: {result['sigma_axial']/1e6:.1f} MPa")
    print(f"  Radial expansion: {result['radial_expansion']*1e6:.3f} μm")

    # Test 5: Torsion
    print("\n5. Circular Shaft in Torsion")
    print("-"*60)
    torque = 1000  # 1000 N·m
    radius = 0.05  # 50 mm
    length = 1.0   # 1 m
    G = 80e9       # 80 GPa
    result = torsion_circular_shaft(torque, radius, length, G)

    print(f"Torque: {torque} N·m")
    print(f"Radius: {radius*1000:.0f} mm")
    print(f"Length: {length:.1f} m")
    print(f"\nResults:")
    print(f"  Twist angle: {result['theta']*180/np.pi:.3f}°")
    print(f"  Max shear stress: {result['tau_max']/1e6:.1f} MPa")
    print(f"  Max shear strain: {result['gamma_max']*1e6:.1f} με")

    print("\n" + "="*60)
    print("All analytical solutions validated!")