"""
Deethanizer Distillation Column Model - Rewritten solver.

Implements the Thiele-Geddes method with theta convergence:
- PR EOS from scratch with kij (not thermo library)
- Constant 5.42 in Kb method (eq. 2.13, 5.36-5.37)
- Diphasic feed split on 2 stages (eq. 5.24)
- Bidirectional energy balance (eq. 5.53-5.67)
- Modified Thomas algorithm (Boston & Sullivan, eq. 5.26)

Reference: Chapitre 5 & 6 - Application au deethaniseur

Stages numbering:
    j=0: Condenser (partial)
    j=1 to N: Trays (N=28)
    j=N+1=29: Reboiler (partial)

Feed convention (Case 1, j_FC=5, j_FH=21):
    j=4: VFC enters (cold feed vapor)
    j=5: LFC enters (cold feed liquid)
    j=20: VFH enters (hot feed vapor)
    j=21: LFH enters (hot feed liquid)
"""

import numpy as np
from .PengRobinson import PengRobinsonEOS, R, COMPONENT_DB


class Deethanizer:
    """Deethanizer column simulation using theta convergence method."""

    COMPONENTS = [
        "C1", "C2", "C3", "iC4", "nC4", "iC5", "nC5",
        "C6", "C7", "C8", "C9", "C10", "C11", "C12",
        "CO2", "N2"
    ]

    def __init__(self, n_trays=28, components=None):
        if components is None:
            components = self.COMPONENTS
        self.components = components
        self.nc = len(components)
        self.N = n_trays
        self.n_stages = n_trays + 2  # 0=condenser .. N+1=reboiler
        self.eos = PengRobinsonEOS(components)
        self.P = 2.74912e6  # Pa (27.49 bar)

        # Feeds
        self.FC_total = 0.0
        self.FH_total = 0.0
        self.FC_z = np.zeros(self.nc)
        self.FH_z = np.zeros(self.nc)
        self.FC_T = 0.0
        self.FH_T = 0.0
        self.FC_P = 0.0
        self.FH_P = 0.0
        self.j_FC = 5   # liquid cold feed stage (vapor enters at j_FC-1)
        self.j_FH = 21  # liquid hot feed stage (vapor enters at j_FH-1)

        # Specs
        self.D_total = 0.0
        self.B_total = 0.0
        self.L0 = 0.0

        # Flash results (component flows in kmol/h)
        self.VFC = np.zeros(self.nc)
        self.LFC = np.zeros(self.nc)
        self.VFH = np.zeros(self.nc)
        self.LFH = np.zeros(self.nc)
        self.VFC_total = 0.0
        self.LFC_total = 0.0
        self.VFH_total = 0.0
        self.LFH_total = 0.0

        # Stage profiles
        self.T = np.zeros(self.n_stages)
        self.V = np.zeros(self.n_stages)
        self.L = np.zeros(self.n_stages)
        self.K = np.zeros((self.nc, self.n_stages))
        self.v = np.zeros((self.nc, self.n_stages))
        self.l = np.zeros((self.nc, self.n_stages))

        # Results
        self.d = np.zeros(self.nc)
        self.b = np.zeros(self.nc)
        self.QC = 0.0
        self.QR = 0.0
        self.converged = False
        self.iterations = 0

    # ================================================================
    # Case setup (data from Chapter 6)
    # ================================================================
    def set_default_case1(self):
        """Case 1: current operation. FC@5, FH@21, L0=306.2."""
        self._set_common_data()
        self.j_FC = 5
        self.j_FH = 21
        self.L0 = 306.2

    def set_case2(self):
        """Case 2: FC@6, FH@19, L0=306.2."""
        self._set_common_data()
        self.j_FC = 6
        self.j_FH = 19
        self.L0 = 306.2

    def set_case3(self):
        """Case 3: FC@6, FH@19, L0=236.1 (optimal reflux)."""
        self._set_common_data()
        self.j_FC = 6
        self.j_FH = 19
        self.L0 = 236.1

    def _set_common_data(self):
        self.FC_total = 923.9056736
        self.FC_T = 273.15 - 18
        self.FC_P = 27.6873e5
        self.FC_z = np.array([
            0.2894, 0.2027, 0.1641, 0.0499, 0.1026,
            0.0352, 0.0488, 0.0525, 0.0224, 0.0130,
            0.0071, 0.0039, 0.0018, 0.0012,
            0.0025, 0.0029
        ])
        self.FH_total = 362.6380771
        self.FH_T = 273.15 + 115
        self.FH_P = 28.3738e5
        self.FH_z = np.array([
            0.1143, 0.0711, 0.0655, 0.0234, 0.0517,
            0.0236, 0.0365, 0.0678, 0.0532, 0.0586,
            0.0603, 0.0616, 0.0529, 0.2568,
            0.0010, 0.0017
        ])
        self.D_total = 562.1
        self.B_total = 724.443

    # ================================================================
    # Reference flash data from Tables 6-3 and 6-4
    # ================================================================
    def _load_reference_flash(self):
        """Use the exact flash results from the reference tables."""
        # Cold feed (Table 6-3): VFCi, LFCi in kmol/h
        self.VFC = np.array([
            136.8441, 24.93438, 5.494693, 0.635093, 0.881636,
            0.114947, 0.112816, 0.0343, 4.34e-3, 7.53e-4,
            1.31e-4, 2.39e-5, 3.46e-6, 8.90e-7,
            0.62983, 2.12036
        ])
        self.LFC = np.array([
            130.534161, 162.341303, 146.118228, 45.4678003, 93.9110864,
            32.4065327, 44.9737804, 48.4707068, 20.6911423, 12.0100204,
            6.55959944, 3.60320827, 1.66302676, 1.10868592,
            1.67993457, 0.55896662
        ])
        self.VFC_total = np.sum(self.VFC)
        self.LFC_total = np.sum(self.LFC)

        # Hot feed (Table 6-4): VFHi, LFHi in kmol/h
        self.VFH = np.array([
            21.4891, 7.44029, 3.85333, 0.852106, 1.586446,
            0.429948, 0.581043, 0.558364, 0.230459, 0.133964,
            0.074024, 0.041135, 0.019089, 0.053132,
            0.142604, 0.419114
        ])
        self.LFH = np.array([
            19.96043, 18.34328, 19.89946, 7.633625, 17.16194,
            8.128311, 12.65525, 24.0285, 19.06189, 21.11663,
            21.79305, 22.29737, 19.16446, 93.07233,
            0.220034, 0.197371
        ])
        self.VFH_total = np.sum(self.VFH)
        self.LFH_total = np.sum(self.LFH)

    # ================================================================
    # Initial profiles from Table 6-6, 6-9, 6-12
    # ================================================================
    def _init_profiles_case1(self):
        """Initial T, V, L from Table 6-6 (Case 1)."""
        self.T = np.array([
            255.8, 275.9, 282.9, 286.8, 290.7, 299.9, 326.7, 339.0,
            345.6, 349.5, 351.9, 353.4, 354.3, 354.9, 355.2, 355.5,
            355.7, 356.0, 356.4, 357.0, 358.6, 369.0, 372.3, 374.0,
            375.5, 377.4, 380.7, 386.8, 400.0, 437.1
        ], dtype=float)
        self.V = np.array([
            562.1, 868.3, 861.5, 849.9, 833.4, 791.4, 812.5, 1033.9,
            1155.0, 1226.8, 1272.6, 1302.3, 1321.3, 1333.4, 1341.0,
            1345.7, 1348.6, 1350.0, 1349.9, 1347.0, 1334.3, 1219.2,
            1139.7, 1216.0, 1244.3, 1256.4, 1261.3, 1256.4, 1216.4,
            1025.4
        ], dtype=float)
        self.L = np.array([
            306.2, 299.4, 287.8, 271.3, 401.1074, 1174.3055, 1395.7055,
            1516.8055, 1588.6055, 1634.4055, 1664.1055, 1683.1055,
            1695.2055, 1702.8055, 1707.5055, 1710.4055, 1711.8055,
            1711.7055, 1708.8055, 1696.1055, 1618.9095, 1864.1434,
            1940.4434, 1968.7334, 1980.8334, 1985.7334, 1980.8334,
            1940.8334, 1749.8334, 724.4434
        ], dtype=float)

    def _init_profiles_case2(self):
        """Initial T, V, L from Table 6-9 (Case 2)."""
        self.T = np.array([
            255.8, 275.9, 282.9, 286.8, 290.7, 299.9, 326.7, 339.0,
            345.6, 349.5, 351.9, 353.4, 354.3, 354.9, 355.2, 355.5,
            355.7, 356.0, 356.4, 357.0, 358.6, 369.0, 372.3, 374.0,
            375.5, 377.4, 380.7, 386.8, 400.0, 437.1
        ], dtype=float)
        self.V = np.array([
            562.1, 868.3, 861.5, 849.9, 833.4, 791.4, 812.5, 1033.9,
            1155.0, 1226.8, 1272.6, 1302.3, 1321.3, 1333.4, 1341.0,
            1345.7, 1348.6, 1350.0, 1349.9, 1347.0, 1334.3, 1219.2,
            1139.7, 1216.0, 1244.3, 1256.4, 1261.3, 1256.4, 1216.4,
            1025.4
        ], dtype=float)
        self.L = np.array([
            306.2, 299.4, 287.8, 271.3, 229.3, 422.2074, 1395.7055,
            1516.8055, 1588.6055, 1634.4055, 1664.1055, 1683.1055,
            1695.2055, 1702.8055, 1707.5055, 1710.4055, 1711.8055,
            1711.7055, 1746.7095, 2058.7434, 1943.6434, 1864.1434,
            1940.4434, 1968.7334, 1980.8334, 1985.7334, 1980.8334,
            1940.8334, 1749.8334, 724.4434
        ], dtype=float)

    def _init_profiles_case3(self):
        """Initial T, V, L from Table 6-12 (Case 3)."""
        self.T = np.array([
            255.8, 275.9, 282.9, 286.8, 290.7, 299.9, 326.7, 339.0,
            345.6, 349.5, 351.9, 353.4, 354.3, 354.9, 355.2, 355.5,
            355.7, 356.0, 356.4, 357.0, 358.6, 369.0, 372.3, 374.0,
            375.5, 377.4, 380.7, 386.8, 400.0, 437.1
        ], dtype=float)
        self.V = np.array([
            562.1, 798.2, 861.5, 849.9, 833.4, 771.4, 812.5, 1033.9,
            1155.0, 1226.8, 1272.6, 1302.3, 1321.3, 1333.4, 1341.0,
            1345.7, 1348.6, 1350.0, 1349.9, 1347.0, 1334.3, 1219.2,
            1139.7, 1216.0, 1244.3, 1256.4, 1261.3, 1256.4, 1216.4,
            1025.4
        ], dtype=float)
        self.L = np.array([
            236.1, 299.4, 287.8, 271.3, 229.3, 442.2074, 1415.7055,
            1516.8055, 1588.6055, 1634.4055, 1664.1055, 1683.1055,
            1695.2055, 1702.8055, 1707.5055, 1710.4055, 1711.8055,
            1711.7055, 1746.7095, 2058.7434, 1943.6434, 1864.1434,
            1940.4434, 1968.7334, 1980.8334, 1985.7334, 1980.8334,
            1940.8334, 1749.8334, 724.4434
        ], dtype=float)

    # ================================================================
    # K-values
    # ================================================================
    def _compute_K_wilson(self):
        """Wilson K-values at each stage."""
        for j in range(self.n_stages):
            self.K[:, j] = self.eos.K_values_wilson(self.T[j], self.P)

    def _compute_K_542(self):
        """
        K-values using eq. 2.13 for ALL components (not just reference).
        K_ji = (Pc_i / P) * exp(5.42 * (1 - Tc_i / T_j))

        This is the correlation from the document, using the same constant
        5.42 for all components (no omega dependency).
        """
        Tc = self.eos.Tc
        Pc = self.eos.Pc
        P = self.P
        for j in range(self.n_stages):
            self.K[:, j] = (Pc / P) * np.exp(5.42 * (1.0 - Tc / self.T[j]))

    def _compute_K_PR(self):
        """PR K-values from scratch (phi_L/phi_V) at each stage."""
        for j in range(self.n_stages):
            x_j = np.maximum(self.l[:, j], 1e-30)
            x_j = x_j / np.sum(x_j)
            y_j = np.maximum(self.v[:, j], 1e-30)
            y_j = y_j / np.sum(y_j)
            try:
                K_new = self.eos.K_values(self.T[j], self.P, x_j, y_j)
                if np.all(K_new > 0) and np.all(np.isfinite(K_new)):
                    self.K[:, j] = K_new
            except Exception:
                pass  # keep previous K

    # ================================================================
    # Modified Thomas algorithm with feed split (eq. 5.24, 5.26)
    # ================================================================
    def _solve_tridiagonal(self, i):
        """
        Solve for liquid component flows l_ji using modified Thomas
        algorithm (Boston & Sullivan, eq. 5.26) with diphasic feed
        split (eq. 5.24).

        Feed split convention:
            VFC enters at j_FC - 1, LFC enters at j_FC
            VFH enters at j_FH - 1, LFH enters at j_FH
        """
        NS = self.n_stages

        # Stripping factor S_ji = K_ji * V_j / L_j
        S = np.zeros(NS)
        for j in range(NS):
            S[j] = self.K[i, j] * self.V[j] / max(self.L[j], 1e-10)

        # Feed vector with diphasic split (eq. 5.24)
        F = np.zeros(NS)
        F[self.j_FC - 1] += self.VFC[i]   # vapor cold feed at j_FC-1
        F[self.j_FC]     += self.LFC[i]   # liquid cold feed at j_FC
        F[self.j_FH - 1] += self.VFH[i]   # vapor hot feed at j_FH-1
        F[self.j_FH]     += self.LFH[i]   # liquid hot feed at j_FH

        n = NS

        # Modified Thomas algorithm (eq. 5.26)
        bpp = np.zeros(n)
        bp = np.zeros(n)
        fp = np.zeros(n)
        cp = np.zeros(n)

        # Forward sweep
        bpp[0] = S[0]
        bp[0] = 1.0 + bpp[0]
        fp[0] = F[0] / bp[0]

        for k in range(1, n):
            bpp[k] = S[k] * bpp[k - 1] / bp[k - 1]
            bp[k] = 1.0 + bpp[k]
            cp[k - 1] = S[k] / bp[k - 1]
            fp[k] = (F[k] + fp[k - 1]) / bp[k]

        # Back-substitution
        l_i = np.zeros(n)
        l_i[n - 1] = fp[n - 1]
        for k in range(n - 2, -1, -1):
            l_i[k] = fp[k] + cp[k] * l_i[k + 1]

        return np.maximum(l_i, 0.0)

    # ================================================================
    # Theta convergence (eq. 5.27-5.33)
    # ================================================================
    def _theta_convergence(self):
        """
        Find theta such that sum(d_i_corrected) = D_total.

        eq. 5.30: d_i_co = (FC_i+FH_i) / (1 + theta*(b_i/d_i)_ca)
        eq. 5.31: g(theta) = sum(d_i_co) - D = 0
        eq. 5.32: g'(theta) for Newton-Raphson
        eq. 5.33: b_i_co = d_i_co * theta * (b_i/d_i)_ca
        """
        NS = self.n_stages

        # Calculated d and b from Thomas solve
        d_calc = np.zeros(self.nc)
        b_calc = np.zeros(self.nc)
        for i in range(self.nc):
            S0 = self.K[i, 0] * self.V[0] / max(self.L[0], 1e-10)
            d_calc[i] = S0 * self.l[i, 0]  # eq. 5.13: d_i = S_0i * l_0i
            b_calc[i] = self.l[i, NS - 1]   # b_i = l_{N+1,i}

        # Total feed per component
        f_i = self.VFC + self.LFC + self.VFH + self.LFH

        # Newton-Raphson for theta (eq. 5.31-5.32), theta_init = 1
        theta = 1.0
        for _ in range(200):
            g = -self.D_total
            dg = 0.0
            for i in range(self.nc):
                denom = d_calc[i] + theta * b_calc[i]
                if denom > 1e-30:
                    g += f_i[i] * d_calc[i] / denom
                    dg -= f_i[i] * d_calc[i] * b_calc[i] / denom ** 2
            if abs(dg) < 1e-30:
                break
            delta = -g / dg
            theta_new = theta + delta
            if theta_new <= 0:
                theta_new = theta * 0.5
            if abs(delta) < 1e-12 * abs(theta):
                theta = theta_new
                break
            theta = theta_new

        # Corrected flows (eq. 5.30, 5.33)
        for i in range(self.nc):
            denom = d_calc[i] + theta * b_calc[i]
            if denom > 1e-30:
                self.d[i] = f_i[i] * d_calc[i] / denom
                self.b[i] = f_i[i] - self.d[i]
            else:
                self.d[i] = 0.0
                self.b[i] = f_i[i]

        # Scale internal flows (eq. 5.34-5.35)
        for i in range(self.nc):
            total_calc = d_calc[i] + b_calc[i]
            if total_calc > 1e-30:
                scale = f_i[i] / total_calc
            else:
                scale = 1.0
            self.l[i, :] *= scale
            self.v[i, :] *= scale

    # ================================================================
    # Temperature update: Kb method with constant 5.42 (eq. 2.13, 5.36-5.37)
    # ================================================================
    def _update_temperatures_Kb_542(self):
        """
        Kb method using constant 5.42 from eq. 2.13.

        eq. 5.36: K_bj_new = 1 / sum(alpha_ji * x_ji)
        eq. 2.13: K_jb = (Pc_b/P) * exp(5.42 * (1 - Tc_b/T))
        eq. 5.37: T_{j,n+1} = 5.42*Tc_b / (5.42 + ln(Pc_b/(K_jb*P)))

        Reference component b = C3 (propane), index 2.
        """
        Tc = self.eos.Tc
        Pc = self.eos.Pc
        P = self.P
        b_ref = 2  # C3 = propane (median component)
        C = 5.42   # Constant from eq. 2.13

        for j in range(self.n_stages):
            x_j = np.maximum(self.l[:, j], 1e-30)
            x_j = x_j / np.sum(x_j)

            K_j = self.K[:, j]
            K_bj = K_j[b_ref]
            if K_bj < 1e-20:
                continue

            # Relative volatilities (eq. 5.36)
            alpha_j = K_j / K_bj
            sum_alpha_x = np.dot(alpha_j, x_j)
            if sum_alpha_x < 1e-20:
                continue

            # New Kb (eq. 5.36)
            K_bj_new = 1.0 / sum_alpha_x

            # Temperature inversion (eq. 5.37)
            ratio = Pc[b_ref] / (K_bj_new * P)
            if ratio <= 0:
                continue
            denom = C + np.log(ratio)
            if denom <= 0.01:
                continue

            T_new = C * Tc[b_ref] / denom
            T_new = np.clip(T_new, 150, 600)
            self.T[j] = T_new

    # ================================================================
    # Bidirectional energy balance (eq. 5.53-5.67)
    # ================================================================
    def _update_flows_energy_balance(self, damping=0.5):
        """
        Bidirectional energy balance with diphasic feed split.

        Top section (eq. 5.53-5.65): L computed, V from material balance
        Bottom section (eq. 5.66-5.67): V computed, L from material balance

        IMPORTANT: For each j from 5 to 21, V_j is computed BEFORE L_j
        (remark from document).
        """
        NS = self.n_stages
        jFC = self.j_FC      # liquid cold feed stage
        jFH = self.j_FH      # liquid hot feed stage
        jV_FC = jFC - 1      # vapor cold feed stage
        jV_FH = jFH - 1      # vapor hot feed stage

        # --- Mixture enthalpies at each stage ---
        h = np.zeros(NS)  # liquid enthalpy (J/mol)
        H = np.zeros(NS)  # vapor enthalpy (J/mol)
        for j in range(NS):
            x_j = np.maximum(self.l[:, j], 1e-30)
            x_j = x_j / np.sum(x_j)
            y_j = np.maximum(self.v[:, j], 1e-30)
            y_j = y_j / np.sum(y_j)
            h[j] = self.eos.pure_component_enthalpy(self.T[j], self.P, x_j, "liquid")
            H[j] = self.eos.mixture_enthalpy(self.T[j], self.P, y_j, "vapor")

        # Feed enthalpies
        xFC = self.LFC / max(self.LFC_total, 1e-30)
        yFC = self.VFC / max(self.VFC_total, 1e-30)
        xFH = self.LFH / max(self.LFH_total, 1e-30)
        yFH = self.VFH / max(self.VFH_total, 1e-30)
        h_FC = self.eos.pure_component_enthalpy(self.FC_T, self.FC_P, xFC, "liquid")
        H_FC = self.eos.mixture_enthalpy(self.FC_T, self.FC_P, yFC, "vapor")
        h_FH = self.eos.pure_component_enthalpy(self.FH_T, self.FH_P, xFH, "liquid")
        H_FH = self.eos.mixture_enthalpy(self.FH_T, self.FH_P, yFH, "vapor")

        D = self.D_total
        B = self.B_total
        HD = H[0]           # H_D = vapor enthalpy at condenser
        hB = h[NS - 1]      # h_B = liquid enthalpy at reboiler

        V_new = np.zeros(NS)
        L_new = np.zeros(NS)

        # === TOP SECTION ===

        # (5.53) V[1] = L[0] + D
        L_new[0] = self.L0
        V_new[0] = D
        V_new[1] = self.L0 + D

        # (5.54) QC = V[1]*H[1] - L[0]*h[0] - D*H_D
        QC = V_new[1] * H[1] - L_new[0] * h[0] - D * HD

        # (5.55-5.56) No-feed stages: j=1 to jV_FC-1
        for j in range(1, jV_FC):
            dH = H[j + 1] - h[j]
            if abs(dH) > 1e-3:
                L_new[j] = (QC + D * (HD - H[j + 1])) / dH
            else:
                L_new[j] = self.L[j]
            L_new[j] = max(L_new[j], 1.0)
            V_new[j + 1] = L_new[j] + D  # (5.56)

        # (5.57) Stage jV_FC: VFC enters (vapor cold feed)
        j = jV_FC
        dH = H[j + 1] - h[j]
        if abs(dH) > 1e-3:
            L_new[j] = (QC + D * (HD - H[j + 1])
                        + self.VFC_total * (H[j + 1] - H_FC)) / dH
        else:
            L_new[j] = self.L[j]
        L_new[j] = max(L_new[j], 1.0)

        # (5.58) V[jFC] = -VFC + L[jV_FC] + D
        V_new[j + 1] = -self.VFC_total + L_new[j] + D

        # (5.59) Stage jFC: LFC enters (liquid cold feed)
        j = jFC
        dH = H[j + 1] - h[j]
        if abs(dH) > 1e-3:
            L_new[j] = (self.LFC_total * (H[j + 1] - h_FC)
                        + L_new[j - 1] * (H[j + 1] - h[j - 1])
                        + V_new[j] * (H[j] - H[j + 1])) / dH
        else:
            L_new[j] = self.L[j]
        L_new[j] = max(L_new[j], 1.0)

        # (5.60) V[jFC+1] = -LFC - L[jV_FC] + V[jFC] + L[jFC]
        V_new[j + 1] = -self.LFC_total - L_new[j - 1] + V_new[j] + L_new[j]

        # (5.61-5.62) Middle section: j=jFC+1 to jV_FH-1
        for j in range(jFC + 1, jV_FH):
            dH = H[j + 1] - h[j]
            if abs(dH) > 1e-3:
                L_new[j] = (L_new[j - 1] * (H[j + 1] - h[j - 1])
                            + V_new[j] * (H[j] - H[j + 1])) / dH
            else:
                L_new[j] = self.L[j]
            L_new[j] = max(L_new[j], 1.0)
            V_new[j + 1] = -L_new[j - 1] + V_new[j] + L_new[j]  # (5.62)

        # (5.63) Stage jV_FH: VFH enters (vapor hot feed)
        j = jV_FH
        dH = H[j + 1] - h[j]
        if abs(dH) > 1e-3:
            L_new[j] = (L_new[j - 1] * (H[j + 1] - h[j - 1])
                        + self.VFH_total * (H[j + 1] - H_FH)
                        + V_new[j] * (H[j] - H[j + 1])) / dH
        else:
            L_new[j] = self.L[j]
        L_new[j] = max(L_new[j], 1.0)

        # (5.64) V[jFH] = -L[jV_FH-1] - VFH + V[jV_FH] + L[jV_FH]
        V_new[j + 1] = -L_new[j - 1] - self.VFH_total + V_new[j] + L_new[j]

        # (5.65) Stage jFH: LFH enters (liquid hot feed)
        j = jFH
        dH = H[j + 1] - h[j]
        if abs(dH) > 1e-3:
            L_new[j] = (L_new[j - 1] * (H[j + 1] - h[j - 1])
                        + self.LFH_total * (H[j + 1] - h_FH)
                        + V_new[j] * (H[j] - H[j + 1])) / dH
        else:
            L_new[j] = self.L[j]
        L_new[j] = max(L_new[j], 1.0)
        V_new[j + 1] = -L_new[j - 1] - self.LFH_total + V_new[j] + L_new[j]

        # === BOTTOM SECTION (eq. 5.66-5.67) ===

        # (5.51) QR from global energy balance
        hF_total = (self.LFC_total * h_FC + self.VFC_total * H_FC
                    + self.LFH_total * h_FH + self.VFH_total * H_FH)
        QR = -hF_total + B * hB + D * HD + QC

        # (5.66) V[j] for j=jFH+2 to N+1
        for j in range(jFH + 2, NS):
            dH = H[j] - h[j - 1]
            if abs(dH) > 1e-3:
                V_new[j] = (QR - B * (hB - h[j - 1])) / dH
            else:
                V_new[j] = self.V[j]
            V_new[j] = max(V_new[j], 1.0)
            # (5.67) L[j-1] = V[j] + B
            L_new[j - 1] = V_new[j] + B

        L_new[NS - 1] = B

        # --- Damped update ---
        self.V = (1 - damping) * self.V + damping * np.maximum(V_new, 1.0)
        self.L = (1 - damping) * self.L + damping * np.maximum(L_new, 1.0)
        self.V[0] = D
        self.L[0] = self.L0
        self.L[NS - 1] = B

        self.QC = QC
        self.QR = QR

    # ================================================================
    # Main solver
    # ================================================================
    def solve(self, case=1, max_outer=20, tol=1.0, verbose=True):
        """
        Run the theta convergence simulation.

        Loop per iteration:
            1. Thomas algorithm with feed split (eq. 5.24, 5.26)
            2. Theta convergence (eq. 5.27-5.33)
            3. T update: Kb with constant 5.42 (eq. 5.36-5.37, 2.13)
            4. K-values: PR from scratch with kij
            5. Energy balance: bidirectional (eq. 5.53-5.67)
        """
        # Setup
        if case == 1:
            self.set_default_case1()
            self._init_profiles_case1()
        elif case == 2:
            self.set_case2()
            self._init_profiles_case2()
        elif case == 3:
            self.set_case3()
            self._init_profiles_case3()

        self._load_reference_flash()

        if verbose:
            print("=" * 60)
            print("Deethanizer - PR+kij / 5.42 / feed split / eq.5.53-5.67")
            print("=" * 60)
            print(f"Case {case}: FC@{self.j_FC}, FH@{self.j_FH}, L0={self.L0}")
            print(f"D={self.D_total}, B={self.B_total} kmol/h")
            print(f"VFC={self.VFC_total:.2f}, LFC={self.LFC_total:.2f}")
            print(f"VFH={self.VFH_total:.2f}, LFH={self.LFH_total:.2f}")
            print(f"kij: {np.count_nonzero(self.eos.kij)} non-zero pairs")
            print("-" * 60)

        # K-values: eq. 2.13 with constant 5.42 for ALL components
        self._compute_K_542()

        for it in range(max_outer):
            T_old = self.T.copy()

            # 1. Thomas algorithm with feed split (eq. 5.24, 5.26)
            for i in range(self.nc):
                self.l[i, :] = self._solve_tridiagonal(i)
                for j in range(self.n_stages):
                    S_ij = self.K[i, j] * self.V[j] / max(self.L[j], 1e-10)
                    self.v[i, j] = S_ij * self.l[i, j]

            # 2. Theta convergence (eq. 5.27-5.33)
            self._theta_convergence()

            # 3. T update: Kb with constant 5.42 (eq. 5.36-5.37)
            self._update_temperatures_Kb_542()

            # 4. K-values: eq. 2.13 at new T (for next iteration)
            self._compute_K_542()

            # 5. Energy balance (eq. 5.53-5.67)
            self._update_flows_energy_balance(damping=0.5)

            # Convergence check
            dT_max = np.max(np.abs(self.T - T_old))
            if verbose:
                print(f"  Iter {it+1}: max|dT|={dT_max:.2f} K, "
                      f"T_top={self.T[0]:.1f} K ({self.T[0]-273.15:.1f} C), "
                      f"T_bot={self.T[-1]:.1f} K ({self.T[-1]-273.15:.1f} C), "
                      f"QC={abs(self.QC):,.0f} QR={abs(self.QR):,.0f} kJ/h")
            if dT_max < tol:
                self.converged = True
                self.iterations = it + 1
                break

        if not self.converged:
            self.iterations = max_outer

        if verbose:
            self._print_results()

        return self._get_results()

    def _print_results(self):
        print("\n" + "=" * 60)
        if self.converged:
            print(f"CONVERGED in {self.iterations} iterations")
        else:
            print(f"NOT converged after {self.iterations} iterations")
        print("=" * 60)

        print(f"\nQC = {abs(self.QC):,.2f} kJ/h")
        print(f"QR = {abs(self.QR):,.2f} kJ/h")

        print(f"\n{'Stage':>5} {'T(K)':>8} {'T(C)':>8} {'V(kmol/h)':>12} {'L(kmol/h)':>12}")
        for j in range(self.n_stages):
            print(f"{j:>5} {self.T[j]:>8.1f} {self.T[j]-273.15:>8.1f} "
                  f"{self.V[j]:>12.1f} {self.L[j]:>12.1f}")

        print(f"\n{'Comp':>5} {'d_i':>14} {'b_i':>14}")
        for i in range(self.nc):
            print(f"{self.components[i]:>5} {self.d[i]:>14.6f} {self.b[i]:>14.6f}")
        print(f"{'Sum':>5} {np.sum(self.d):>14.4f} {np.sum(self.b):>14.4f}")

        print("\n--- Reference (Case 1, Table 6-7) ---")
        print("T_top = 260.2 K (-12.95 C), T_bot = 416.6 K (143.45 C)")
        print("QC = 4,133,587.91 kJ/h, QR = 22,854,385.49 kJ/h")

    def _get_results(self):
        return {
            "converged": self.converged,
            "iterations": self.iterations,
            "T": self.T.copy(),
            "V": self.V.copy(),
            "L": self.L.copy(),
            "K": self.K.copy(),
            "d": self.d.copy(),
            "b": self.b.copy(),
            "QC": self.QC,
            "QR": self.QR,
            "components": self.components,
        }
