# ============================================================
# Moment–Curvature (M–kappa) for RC section (State II cracked)
# - Concrete: tension = 0, compression linear elastic (Ec)
# - Steel: bilinear (elastic-perfectly plastic via stress cap at +/- fyd)
# - Solve eps0 for given kappa and axial force N_Ed by scan + bisection
# - Failure when eps_c reaches -eps_cu OR eps_s reaches +/- eps_su
#
# UNITS:
#   m, MN, MN*m, MPa (= MN/m^2)
# SIGN:
#   eps > 0  tension
#   eps < 0  compression
# GEOMETRY:
#   z positive upwards
# CURVATURE definition:
#   eps(z) = eps0 - kappa * z
#   kappa > 0 => top (z>0) in compression (eps negative)  ✅ sagging
# MOMENT:
#   My = -∫ sigma * z dA  (sagging positive)
# ============================================================

import numpy as np
from dataclasses import dataclass
from typing import Dict, Any, Optional, Tuple

# ----------------------------
# Data structures
# ----------------------------

@dataclass
class ConcreteMesh:
    y: np.ndarray   # [m] not used in 1D bending, but kept
    z: np.ndarray   # [m]
    A: np.ndarray   # [m^2]

@dataclass
class SteelBars:
    y: np.ndarray   # [m] not used in 1D bending, but kept
    z: np.ndarray   # [m]
    As: np.ndarray  # [m^2]

# ----------------------------
# Material laws (state II)
# ----------------------------

class ConcreteCrackedLinear:
    """Concrete: compression linear (eps<0), tension zero (eps>=0). Stress in MPa (neg. in compression)."""
    def __init__(self, Ec_MPa: float):
        self.Ec = float(Ec_MPa)

    def sigma(self, eps: np.ndarray) -> np.ndarray:
        # eps<0: compression => sigma negative; eps>=0: no tension
        return np.where(eps < 0.0, self.Ec * eps, 0.0)

class ConcreteEC2ParabolaRectangle:
    """
    EC2 parabola-rectangle stress-strain law for concrete (compression only).
    Stress in MPa (negative in compression).
    """
    def __init__(self, fcd, eps_c2=2.0e-3, eps_cu=3.5e-3):
        self.fcd = float(fcd)
        self.eps_c2 = float(eps_c2)
        self.eps_cu = float(eps_cu)

    def sigma(self, eps):
        sig = np.zeros_like(eps)

        # compression only
        comp = eps < 0.0
        e = -eps[comp]  # positive compression strain

        # parabola
        mask_p = e <= self.eps_c2
        sig[comp][mask_p] = -self.fcd * (
            1.0 - (1.0 - e[mask_p]/self.eps_c2)**2
        )

        # rectangle
        mask_r = (e > self.eps_c2) & (e <= self.eps_cu)
        sig[comp][mask_r] = -self.fcd

        return sig



class SteelElasticPerfectlyPlastic:
    """
    Steel: linear elastic up to yield, then perfectly plastic (stress capped at +/- fyd).
    Stress in MPa.
    """
    def __init__(self, Es_MPa: float, fyd_MPa: float):
        self.Es = float(Es_MPa)
        self.fyd = float(fyd_MPa)
        self.eps_y = self.fyd / self.Es

    def sigma(self, eps: np.ndarray) -> np.ndarray:
        sig = self.Es * eps
        return np.clip(sig, -self.fyd, self.fyd)

# ----------------------------
# Simple section meshing
# ----------------------------

def mesh_circle(R, nr=80, ntheta=160):
    r = np.linspace(0, R, nr)
    t = np.linspace(0, 2*np.pi, ntheta, endpoint=False)
    dr, dt = r[1]-r[0], t[1]-t[0]
    Rg, Tg = np.meshgrid(r, t, indexing="ij")
    y = Rg*np.cos(Tg)
    z = Rg*np.sin(Tg)
    A = Rg*dr*dt
    mask = Rg > 0
    return ConcreteMesh(y[mask], z[mask], A[mask])

def mesh_I(bf, tf, tw, h, ny=140, nz=160):
    y = np.linspace(-bf/2, bf/2, ny)
    z = np.linspace(-h/2, h/2, nz)
    dy, dz = y[1]-y[0], z[1]-z[0]
    Y, Z = np.meshgrid(y, z, indexing="ij")
    mask = (
        (np.abs(Y) <= tw/2) |
        (Z >= h/2 - tf) |
        (Z <= -h/2 + tf)
    )
    A = np.full(np.count_nonzero(mask), dy*dz)
    return ConcreteMesh(Y[mask], Z[mask], A)

def mesh_T(bf, tf, tw, h, ny=140, nz=160):
    y = np.linspace(-bf/2, bf/2, ny)
    z = np.linspace(0, h, nz)
    dy, dz = y[1]-y[0], z[1]-z[0]
    Y, Z = np.meshgrid(y, z, indexing="ij")
    mask = (Z >= h-tf) | (np.abs(Y) <= tw/2)
    yv, zv = Y[mask], Z[mask]
    A = np.full(len(yv), dy*dz)
    z0 = np.sum(zv*A)/np.sum(A)
    return ConcreteMesh(yv, zv-z0, A)


def mesh_rectangle(b: float, h: float, ny: int = 60, nz: int = 120) -> ConcreteMesh:
    """
    Rectangle concrete mesh.
    b: width in y direction [m]
    h: height in z direction [m]
    """
    y = np.linspace(-b/2, b/2, ny)
    z = np.linspace(-h/2, h/2, nz)
    dy, dz = y[1]-y[0], z[1]-z[0]
    Y, Z = np.meshgrid(y, z, indexing="ij")
    A = np.full(Y.size, dy*dz, dtype=float)
    return ConcreteMesh(Y.ravel(), Z.ravel(), A)

# ----------------------------
# M-kappa solver (1D bending about y)
# ----------------------------

class MomentCurvatureSolver:
    def __init__(
        self,
        conc: ConcreteMesh,
        steel: SteelBars,
        conc_law: ConcreteCrackedLinear,
        steel_law: SteelElasticPerfectlyPlastic,
        eps_cu: float = 3.5e-3,     # concrete ultimate compressive strain (positive value used with minus sign)
        eps_su: float = 25e-3       # steel ultimate strain (tension/compression)
    ):
        self.conc = conc
        self.steel = steel
        self.conc_law = conc_law
        self.steel_law = steel_law

        self.eps_cu = float(eps_cu)
        self.eps_su = float(eps_su)

        # Yield strain from steel law
        self.eps_y = float(self.steel_law.eps_y)

        self.z_top = float(np.max(conc.z))
        self.z_bot = float(np.min(conc.z))

    # ---- kinematics ----
    @staticmethod
    def strains(eps0: float, kappa: float, z: np.ndarray) -> np.ndarray:
        # eps(z)=eps0 - kappa*z  (kappa>0 => top compression)
        return eps0 - kappa * z

    # ---- section resultants ----
    def resultants(self, eps0: float, kappa: float) -> Tuple[float, float]:
        eps_c = self.strains(eps0, kappa, self.conc.z)
        sig_c = self.conc_law.sigma(eps_c)  # MPa

        if self.steel.z.size:
            eps_s = self.strains(eps0, kappa, self.steel.z)
            sig_s = self.steel_law.sigma(eps_s)  # MPa
        else:
            sig_s = np.array([], dtype=float)

        # N [MN]  (MPa*m^2 = MN)
        N = float(np.sum(sig_c * self.conc.A) + np.sum(sig_s * self.steel.As))

        # My [MN*m], sagging positive: My = -∫ sigma*z dA
        My = float(-np.sum(sig_c * self.conc.A * self.conc.z) - np.sum(sig_s * self.steel.As * self.steel.z))

        return N, My

    # ---- solve eps0 for given kappa and axial N ----
    def solve_eps0_for_N(
        self,
        kappa: float,
        N_Ed: float,
        eps0_min: float = -0.05,
        eps0_max: float = +0.05,
        n_scan: int = 300,
        tol_N: float = 1e-8,
        max_iter: int = 80
    ) -> Optional[float]:
        """
        Finds eps0 such that N(eps0,kappa)=N_Ed.
        Robust bracketing via scan + bisection.
        """
        eps0_grid = np.linspace(eps0_min, eps0_max, n_scan)
        g = np.empty_like(eps0_grid)

        for i, e0 in enumerate(eps0_grid):
            N, _ = self.resultants(float(e0), kappa)
            g[i] = N - N_Ed

        # find first bracket where sign changes or hits zero
        idx = np.where(np.sign(g[:-1]) * np.sign(g[1:]) <= 0)[0]
        if idx.size == 0:
            return None

        i0 = int(idx[0])
        a, b = float(eps0_grid[i0]), float(eps0_grid[i0+1])
        fa, fb = float(g[i0]), float(g[i0+1])

        # bisection
        for _ in range(max_iter):
            m = 0.5 * (a + b)
            Nm, _ = self.resultants(m, kappa)
            fm = float(Nm - N_Ed)

            if abs(fm) < tol_N:
                return m

            if np.sign(fa) * np.sign(fm) <= 0:
                b, fb = m, fm
            else:
                a, fa = m, fm

        return 0.5 * (a + b)

    # ---- main M-kappa curve ----
    def moment_curvature(
        self,
        N_Ed: float = 0.0,
        kappa_max: float = 0.30,     # 1/m
        dkappa: float = 5e-4,        # 1/m
        tol_N: float = 1e-8
    ) -> Dict[str, Any]:

        kappas, Mys, eps0s = [], [], []
        epsc_min_list, epss_max_list, epss_min_list = [], [], []

        kappa_y_tens = None
        kappa_y_comp = None
        kappa_fail = None
        fail_mode = None

        # identify "compression reinforcement" as bars with largest z (closest to top)
        z_s = self.steel.z
        if z_s.size:
            z_top_s = np.max(z_s)
            z_bot_s = np.min(z_s)
            comp_ids = np.where(np.isclose(z_s, z_top_s, rtol=0, atol=1e-12))[0]
            tens_ids = np.where(np.isclose(z_s, z_bot_s, rtol=0, atol=1e-12))[0]
        else:
            comp_ids = np.array([], dtype=int)
            tens_ids = np.array([], dtype=int)

        k = 0.0
        while k <= kappa_max + 1e-15:

            eps0 = self.solve_eps0_for_N(kappa=k, N_Ed=N_Ed, tol_N=tol_N)
            if eps0 is None:
                break

            N, My = self.resultants(eps0, k)

            # strains for checks
            eps_c = self.strains(eps0, k, self.conc.z)
            if z_s.size:
                eps_s = self.strains(eps0, k, z_s)
            else:
                eps_s = np.array([0.0])

            eps_c_min = float(np.min(eps_c))   # most compressive (negative)
            eps_s_max = float(np.max(eps_s))   # max tension
            eps_s_min = float(np.min(eps_s))   # max compression (negative)

            kappas.append(k)
            Mys.append(My)
            eps0s.append(eps0)
            epsc_min_list.append(eps_c_min)
            epss_max_list.append(eps_s_max)
            epss_min_list.append(eps_s_min)

            # ---- Events: steel yield (now meaningful because steel is capped) ----
            if kappa_y_comp is None and comp_ids.size:
                eps_comp = float(np.min(eps_s[comp_ids]))  # negative if in compression
                if eps_comp <= -self.eps_y:
                    kappa_y_comp = k

            if kappa_y_tens is None and tens_ids.size:
                eps_tens = float(np.max(eps_s[tens_ids]))  # positive if in tension
                if eps_tens >= self.eps_y:
                    kappa_y_tens = k

            # ---- Failure ----
            if eps_c_min <= -self.eps_cu:
                kappa_fail = k
                fail_mode = "concrete_crush (eps_c_min <= -eps_cu)"
                break

            if abs(eps_s_max) >= self.eps_su or abs(eps_s_min) >= self.eps_su:
                kappa_fail = k
                fail_mode = "steel_ultimate (|eps_s| >= eps_su)"
                break

            k += dkappa

        return {
            "kappa": np.array(kappas),
            "My": np.array(Mys),
            "eps0": np.array(eps0s),
            "eps_c_min": np.array(epsc_min_list),
            "eps_s_max": np.array(epss_max_list),
            "eps_s_min": np.array(epss_min_list),
            "kappa_y_tens": kappa_y_tens,
            "kappa_y_comp": kappa_y_comp,
            "kappa_fail": kappa_fail,
            "fail_mode": fail_mode,
            "eps_y": self.eps_y
        }

    # ---- plot ----
    def plot_M_kappa(self, res: Dict[str, Any], title="M-κ (State II)"):
        import matplotlib.pyplot as plt

        k = res["kappa"]
        M = res["My"]

        plt.figure()
        plt.plot(k, M, lw=2)
        plt.xlabel(r"$\kappa$ [1/m]")
        plt.ylabel(r"$M_y$ [MN·m]")
        plt.title(title)
        plt.grid(True)

        # markers
        if res["kappa_y_tens"] is not None:
            ky = res["kappa_y_tens"]
            Myy = np.interp(ky, k, M)
            plt.scatter([ky], [Myy], marker="o")
            plt.text(ky, Myy, "  yield tension steel", va="bottom")

        if res["kappa_y_comp"] is not None:
            ky = res["kappa_y_comp"]
            Myy = np.interp(ky, k, M)
            plt.scatter([ky], [Myy], marker="o")
            plt.text(ky, Myy, "  yield compression steel", va="bottom")

        if res["kappa_fail"] is not None:
            kf = res["kappa_fail"]
            Myf = np.interp(kf, k, M)
            plt.scatter([kf], [Myf], marker="x")
            plt.text(kf, Myf, f"  fail: {res['fail_mode']}", va="bottom")

        plt.show()

# ----------------------------
# Example
# ----------------------------

# if __name__ == "__main__":

#     # geometry
#     conc = mesh_rectangle(b=0.30, h=0.50, ny=60, nz=140)

#     # example reinforcement:
#     # two bottom bars (tension zone)
#     # two top bars (compression zone)
#     steel = SteelBars(
#         y=np.array([-0.10, 0.10, -0.10, 0.10], dtype=float),
#         z=np.array([-0.20, -0.20, +0.20, +0.20], dtype=float),
#         As=np.array([6.3e-4, 6.3e-4, 2.0e-4, 2.0e-4], dtype=float)
#     )

#     # materials
#     Ec = 33000.0   # MPa
#     Es = 200000.0  # MPa
#     fyd = 435.0    # MPa

#     #conc_law = ConcreteCrackedLinear(Ec_MPa=Ec)
#     conc_law = ConcreteEC2ParabolaRectangle(
#         fcd = 0.85 * 30.0 / 1.5,   # Beispiel C30/37
#         eps_c2 = 2.0e-3,
#         eps_cu = 3.5e-3
#     )
#     steel_law = SteelElasticPerfectlyPlastic(Es_MPa=Es, fyd_MPa=fyd)

#     solver = MomentCurvatureSolver(
#         conc=conc,
#         steel=steel,
#         conc_law=conc_law,
#         steel_law=steel_law,
#         eps_cu=3.5e-3,
#         eps_su=25e-3
#     )

#     res = solver.moment_curvature(
#         N_Ed=0.0,
#         kappa_max=0.20,
#         dkappa=5e-4,
#         tol_N=1e-8
#     )

#     print("eps_y        =", res["eps_y"])
#     print("kappa_y_tens =", res["kappa_y_tens"])
#     print("kappa_y_comp =", res["kappa_y_comp"])
#     print("kappa_fail   =", res["kappa_fail"])
#     print("fail_mode    =", res["fail_mode"])

#     solver.plot_M_kappa(res, title="Moment–Krümmung (Zustand II, Stahl elast.-plast., Beton Druck linear)")
