import numpy as np

# ---------- Phase II-B: dimer verification ----------
def dimer_hamiltonian(gamma, J):
    return np.array([[1j * gamma, J], [J, -1j * gamma]], dtype=complex)


def biorthogonal_expectation(H, G):
    """Ground-state biorthogonal expectation for a real-spectrum, nondegenerate H.

    The left ket is selected by matching its H† eigenvalue to the complex
    conjugate of the selected right eigenvalue. This is safer than sorting
    the two spectra independently by their real parts.
    """
    w_right, v_right = np.linalg.eig(H)
    idx_right = np.argmin(w_right.real)
    energy_right = w_right[idx_right]
    psi_R = v_right[:, idx_right]

    w_left, v_left = np.linalg.eig(H.conj().T)
    idx_left = np.argmin(np.abs(w_left - energy_right.conjugate()))
    psi_L = v_left[:, idx_left]

    denominator = psi_L.conj() @ psi_R
    if abs(denominator) < 1e-12:
        raise ZeroDivisionError("Left/right eigenvectors are self-orthogonal or nearly so (near an EP).")
    return (psi_L.conj() @ G @ psi_R) / denominator


def phase_ii_b(J=1.0, gamma=0.5, G=None):
    if G is None:
        G = np.array([[1, 0], [0, -1]], dtype=complex)  # sigma_z
    return biorthogonal_expectation(dimer_hamiltonian(gamma, J), G)


# ---------- Phase II-C: XXZ chain at N=4 ----------
def xxz_chain_N4(gamma, J=1.0, Delta=0.5, b=None):
    # Spin-1/2 operators
    I2 = np.eye(2)
    sx = np.array([[0, 1], [1, 0]], dtype=complex) / 2
    sy = np.array([[0, -1j], [1j, 0]], dtype=complex) / 2
    sz = np.array([[1, 0], [0, -1]], dtype=complex) / 2

    def op(local_operator, site, N=4):
        mats = [I2] * N
        mats[site] = local_operator
        out = mats[0]
        for matrix in mats[1:]:
            out = np.kron(out, matrix)
        return out

    if b is None:
        b = J

    H = np.zeros((16, 16), dtype=complex)
    for j in range(3):
        H += b * (op(sx, j) @ op(sx, j + 1) + op(sy, j) @ op(sy, j + 1))
        H += J * Delta * op(sz, j) @ op(sz, j + 1)

    # Staggered imaginary field; open boundaries, site indices 0..3.
    for j in range(4):
        H += 1j * gamma * ((-1) ** j) * op(sz, j)
    return H


def magnetization_sector_indices(N=4):
    """Return computational-basis indices grouped by twice total Sz."""
    sectors = {}
    for basis_index in range(2**N):
        # In this Kronecker basis, each set bit is a spin-down (Sz=-1/2).
        twice_sz = N - 2 * basis_index.bit_count()
        sectors.setdefault(twice_sz, []).append(basis_index)
    return sectors


def characteristic_discriminant(matrix):
    """Compute the characteristic-polynomial discriminant in double precision.

    For a real-coefficient polynomial, the discriminant is positive when all
    roots are real and distinct; it changes sign at a generic single-pair EP.
    The Sylvester resultant avoids deciding whether a tiny computed imaginary
    part is physical, as in a threshold test on eigvals.
    """
    coeffs_complex = np.poly(matrix)
    if np.max(np.abs(coeffs_complex.imag)) > 1e-8 * max(1.0, np.max(np.abs(coeffs_complex.real))):
        raise ArithmeticError("Sector characteristic polynomial is not numerically real; check PT symmetry.")
    coeffs = coeffs_complex.real
    degree = len(coeffs) - 1
    derivative = np.polyder(coeffs)

    # Sylvester matrix for resultant(p, p'), with descending coefficients.
    sylvester = np.zeros((2 * degree - 1, 2 * degree - 1), dtype=float)
    for row in range(degree - 1):
        sylvester[row, row : row + degree + 1] = coeffs
    for row in range(degree):
        start = degree - 1 + row
        sylvester[start, row : row + degree] = derivative

    resultant = np.linalg.det(sylvester)
    return ((-1) ** (degree * (degree - 1) // 2)) * resultant


def _sector_discriminants(gamma, J=1.0, Delta=0.5, b=None):
    H = xxz_chain_N4(gamma, J=J, Delta=Delta, b=b)
    values = {}
    for twice_sz, indices in magnetization_sector_indices(4).items():
        block = H[np.ix_(indices, indices)]
        values[twice_sz] = characteristic_discriminant(block)
    return values


def find_ep_N4(
    J=1.0,
    Delta=0.5,
    b=None,
    gamma_max=3.0,
    scan_points=3001,
    xtol=1e-12,
    return_details=False,
):
    """Locate the first EP by scanning for and bisecting sector discriminants.

    The coarse scan only finds a sign-change bracket. Bisection then refines
    the earliest positive-to-nonpositive discriminant crossing across all
    conserved magnetization sectors. For a generic second-order EP, the
    discriminant changes sign as one real eigenvalue pair becomes complex.

    Returns the midpoint gamma_EP by default. With return_details=True,
    returns (gamma_EP, twice_Sz_sector, (gamma_lo, gamma_hi)).
    """
    if gamma_max <= 0 or scan_points < 2 or xtol <= 0:
        raise ValueError("Require gamma_max > 0, scan_points >= 2, and xtol > 0.")

    grid = np.linspace(0.0, gamma_max, scan_points)
    previous = _sector_discriminants(grid[0], J=J, Delta=Delta, b=b)
    brackets = []

    for gamma in grid[1:]:
        current = _sector_discriminants(gamma, J=J, Delta=Delta, b=b)
        for twice_sz, d_now in current.items():
            d_prev = previous[twice_sz]
            if d_prev > 0.0 and d_now <= 0.0:
                brackets.append((gamma - grid[1], gamma, twice_sz))
        if brackets:
            # Scanning is ascending; finish this grid point and take the
            # earliest bracket among any sectors that crossed in this interval.
            break
        previous = current

    if not brackets:
        raise RuntimeError(f"No first EP found in 0 <= gamma <= {gamma_max}; increase gamma_max or scan_points.")

    lo, hi, sector = min(brackets, key=lambda item: item[0])
    d_lo = _sector_discriminants(lo, J=J, Delta=Delta, b=b)[sector]
    d_hi = _sector_discriminants(hi, J=J, Delta=Delta, b=b)[sector]
    if not (d_lo > 0.0 and d_hi <= 0.0):
        raise ArithmeticError("Failed to establish a positive-to-nonpositive discriminant bracket.")

    while hi - lo > xtol:
        mid = 0.5 * (lo + hi)
        d_mid = _sector_discriminants(mid, J=J, Delta=Delta, b=b)[sector]
        if d_mid > 0.0:
            lo = mid
        else:
            hi = mid

    gamma_ep = 0.5 * (lo + hi)
    if return_details:
        return gamma_ep, sector, (lo, hi)
    return gamma_ep


def half_filled_sector_gap_N4(gamma, J=1.0, Delta=0.5, b=None):
    """Gap between the two lowest real eigenvalues in the Sz=0 sector."""
    H = xxz_chain_N4(gamma, J=J, Delta=Delta, b=b)
    indices = magnetization_sector_indices(4)[0]
    block = H[np.ix_(indices, indices)]
    eigenvalues = np.linalg.eigvals(block)
    if np.max(np.abs(eigenvalues.imag)) > 1e-8:
        raise ValueError("Requested gap is not below the EP: this sector has complex eigenvalues.")
    energies = np.sort(eigenvalues.real)
    return energies[1] - energies[0]


def puiseux_exponent(gamma_EP, J=1.0, Delta=0.5, b=None, n_points=20):
    """Fit the below-EP Sz=0 coalescing-pair gap against gamma_EP-gamma."""
    eps_vals = np.logspace(-7, -3, n_points)
    gaps = np.array([
        half_filled_sector_gap_N4(gamma_EP - eps, J=J, Delta=Delta, b=b)
        for eps in eps_vals
    ])
    if np.any(gaps <= 0) or not np.all(np.isfinite(gaps)):
        raise ArithmeticError("Nonpositive or nonfinite sector gap; refine EP or adjust the fit window.")
    slope, intercept = np.polyfit(np.log(eps_vals), np.log(gaps), 1)
    return slope, eps_vals, gaps


def _sector_charpoly_gamma_derivatives(gamma, J=1.0, Delta=0.5, b=None, twice_sz=0, step=1e-2):
    """Return characteristic-polynomial coefficients and first two gamma derivatives.

    In a fixed N=4 magnetization sector, each characteristic-polynomial
    coefficient is a polynomial of degree at most the sector dimension in
    gamma. Interpolation on a scaled interval therefore gives stable local
    derivatives without finite-differencing individual eigenvalues.
    """
    indices = magnetization_sector_indices(4)[twice_sz]
    nodes = np.linspace(-1.0, 1.0, 13)
    samples = []
    for node in nodes:
        H = xxz_chain_N4(gamma + step * node, J=J, Delta=Delta, b=b)
        block = H[np.ix_(indices, indices)]
        coeffs = np.poly(block)
        if np.max(np.abs(coeffs.imag)) > 1e-8 * max(1.0, np.max(np.abs(coeffs.real))):
            raise ArithmeticError("Non-real sector characteristic polynomial; check model symmetry.")
        samples.append(coeffs.real)
    samples = np.asarray(samples)
    degree = len(samples[0]) - 1
    # Fit each coefficient against the normalized coordinate x=(gamma'-gamma)/step.
    fits = np.asarray([
        np.polynomial.polynomial.polyfit(nodes, samples[:, k], degree)
        for k in range(samples.shape[1])
    ])
    coeffs = fits[:, 0]
    coeffs_gamma = fits[:, 1] / step
    coeffs_gamma2 = 2.0 * fits[:, 2] / (step * step)
    return coeffs, coeffs_gamma, coeffs_gamma2


def _poly_value(coeffs, energy, derivative_order=0):
    """Evaluate a descending-order polynomial or its energy derivative."""
    derivative = np.polyder(coeffs, derivative_order) if derivative_order else coeffs
    return float(np.polyval(derivative, energy))


def refine_ep_double_root(gamma_seed, J=1.0, Delta=0.5, b=None, twice_sz=0, max_iter=8):
    """Refine a bisection estimate by solving p(E,gamma)=p_E(E,gamma)=0.

    The selected sector's characteristic polynomial has a double root at a
    generic second-order EP. Newton refinement solves the double-root
    equations after bisection has supplied a reliable basin and bracket.
    Returns (gamma_EP, E_EP, residual_norm).
    """
    indices = magnetization_sector_indices(4)[twice_sz]
    H = xxz_chain_N4(gamma_seed, J=J, Delta=Delta, b=b)
    block = H[np.ix_(indices, indices)]
    eigs = np.linalg.eigvals(block)
    _, i, j = min(
        (abs(eigs[i] - eigs[j]), i, j)
        for i in range(len(eigs)) for j in range(i + 1, len(eigs))
    )
    energy = float(((eigs[i] + eigs[j]) / 2.0).real)
    gamma = float(gamma_seed)

    for _ in range(max_iter):
        coeffs, coeffs_g, _ = _sector_charpoly_gamma_derivatives(
            gamma, J=J, Delta=Delta, b=b, twice_sz=twice_sz
        )
        p = _poly_value(coeffs, energy)
        p_e = _poly_value(coeffs, energy, 1)
        p_ee = _poly_value(coeffs, energy, 2)
        p_g = _poly_value(coeffs_g, energy)
        p_eg = _poly_value(coeffs_g, energy, 1)
        jacobian = np.array([[p_e, p_g], [p_ee, p_eg]], dtype=float)
        correction = np.linalg.solve(jacobian, np.array([p, p_e], dtype=float))
        energy -= correction[0]
        gamma -= correction[1]
        if np.linalg.norm(correction, ord=np.inf) < 2e-14:
            break

    coeffs, coeffs_g, _ = _sector_charpoly_gamma_derivatives(
        gamma, J=J, Delta=Delta, b=b, twice_sz=twice_sz
    )
    residual = float(np.linalg.norm([
        _poly_value(coeffs, energy),
        _poly_value(coeffs, energy, 1),
    ], ord=np.inf))
    return gamma, energy, residual


def puiseux_coefficients_N4(gamma_seed, J=1.0, Delta=0.5, b=None, twice_sz=0):
    """Calculate leading Puiseux coefficients from the local characteristic polynomial.

    For delta=gamma-gamma_EP>0, the returned coefficients describe the pair
      E_± = E_EP + c_1 delta + c_2 delta^2
            ± (c_1/2 delta^(1/2) + c_3/2 delta^(3/2)) + O(delta^(5/2)).
    The plus branch is chosen to have positive imaginary part just above the
    PT-breaking threshold. Coefficients follow from a formal series solution
    of p(E,gamma)=0 around its double root, not from fitting eigval gaps.
    """
    gamma_ep, energy_ep, residual = refine_ep_double_root(
        gamma_seed, J=J, Delta=Delta, b=b, twice_sz=twice_sz
    )
    coeffs, coeffs_g, coeffs_gg = _sector_charpoly_gamma_derivatives(
        gamma_ep, J=J, Delta=Delta, b=b, twice_sz=twice_sz
    )
    p20 = _poly_value(coeffs, energy_ep, 2) / 2.0
    p01 = _poly_value(coeffs_g, energy_ep)
    p30 = _poly_value(coeffs, energy_ep, 3) / 6.0
    p11 = _poly_value(coeffs_g, energy_ep, 1)
    p40 = _poly_value(coeffs, energy_ep, 4) / 24.0
    p21 = _poly_value(coeffs_g, energy_ep, 2) / 2.0
    p02 = _poly_value(coeffs_gg, energy_ep) / 2.0
    p50 = _poly_value(coeffs, energy_ep, 5) / 120.0
    p31 = _poly_value(coeffs_g, energy_ep, 3) / 6.0
    p12 = _poly_value(coeffs_gg, energy_ep, 1) / 2.0

    c_half = np.sqrt(complex(-p01 / p20))
    if c_half.imag < 0 or (abs(c_half.imag) < 1e-14 and c_half.real < 0):
        c_half = -c_half
    c_one = -(p30 * c_half**2 + p11) / (2.0 * p20)
    c_three_half = -(
        p20 * c_one**2
        + 3.0 * p30 * c_half**2 * c_one
        + p40 * c_half**4
        + p11 * c_one
        + p21 * c_half**2
        + p02
    ) / (2.0 * p20 * c_half)
    c_two = -(
        2.0 * p20 * c_one * c_three_half
        + 3.0 * p30 * c_half**2 * c_three_half
        + 3.0 * p30 * c_half * c_one**2
        + 4.0 * p40 * c_half**3 * c_one
        + p50 * c_half**5
        + p11 * c_three_half
        + 2.0 * p21 * c_half * c_one
        + p31 * c_half**3
        + p12 * c_half
    ) / (2.0 * p20 * c_half)

    values = {
        "c_1/2": complex(c_half),
        "c_1": complex(c_one),
        "c_3/2": complex(c_three_half),
        "c_2": complex(c_two),
    }
    return gamma_ep, energy_ep, values, residual


def jordan_form_N4(gamma_ep, energy_ep, J=1.0, Delta=0.5, b=None, twice_sz_ep=0):
    """Construct the numerical Jordan canonical form and report block diagnostics.

    The EP sector contributes one size-2 Jordan block. All other roots are
    represented by 1x1 blocks. Repeated eigenvalues in distinct magnetization
    sectors remain separate 1x1 blocks (they are not extra defective blocks).
    Returns (J, block_descriptions, nullities_at_EP).
    """
    H = xxz_chain_N4(gamma_ep, J=J, Delta=Delta, b=b)
    sectors = magnetization_sector_indices(4)
    ep_indices = sectors[twice_sz_ep]
    ep_block = H[np.ix_(ep_indices, ep_indices)]
    ep_eigs = np.linalg.eigvals(ep_block)
    pair = sorted(
        ((abs(ep_eigs[i] - energy_ep) + abs(ep_eigs[j] - energy_ep), i, j)
         for i in range(len(ep_eigs)) for j in range(i + 1, len(ep_eigs)))
    )[0]
    _, i, j = pair
    remaining = [ep_eigs[k] for k in range(len(ep_eigs)) if k not in (i, j)]

    blocks = [np.array([[energy_ep, 1.0], [0.0, energy_ep]], dtype=complex)]
    descriptions = [(2, complex(energy_ep), twice_sz_ep)]
    for value in sorted(remaining, key=lambda z: (z.real, z.imag)):
        blocks.append(np.array([[value]], dtype=complex))
        descriptions.append((1, complex(value), twice_sz_ep))

    for twice_sz in sorted(sectors):
        if twice_sz == twice_sz_ep:
            continue
        indices = sectors[twice_sz]
        block = H[np.ix_(indices, indices)]
        for value in sorted(np.linalg.eigvals(block), key=lambda z: (z.real, z.imag)):
            blocks.append(np.array([[value]], dtype=complex))
            descriptions.append((1, complex(value), twice_sz))

    dimension = sum(block.shape[0] for block in blocks)
    J_form = np.zeros((dimension, dimension), dtype=complex)
    offset = 0
    for block in blocks:
        size = block.shape[0]
        J_form[offset : offset + size, offset : offset + size] = block
        offset += size

    shifted = ep_block - energy_ep * np.eye(len(ep_indices))
    singular_values = np.linalg.svd(shifted, compute_uv=False)
    singular_values_squared = np.linalg.svd(shifted @ shifted, compute_uv=False)
    scale = max(1.0, singular_values[0])
    tolerance = 1e-7 * scale
    nullity_first = int(np.sum(singular_values < tolerance))
    nullity_second = int(np.sum(singular_values_squared < tolerance))
    return J_form, descriptions, (nullity_first, nullity_second)


if __name__ == "__main__":
    print("=== Phase II-B: dimer biorthogonal <sigma_z> ===")
    for gamma in [0.1, 0.5, 0.9]:
        print(f"gamma={gamma}: <sigma_z>_B = {phase_ii_b(gamma=gamma)}")

    print("\n=== Phase II-C: first EP at N=4, via sector discriminant bisection ===")
    gamma_bisected, twice_sz, bracket = find_ep_N4(return_details=True)
    gamma_ep, energy_ep, coefficients, residual = puiseux_coefficients_N4(
        gamma_bisected, twice_sz=twice_sz
    )
    print(f"gamma_EP(N=4) = {gamma_ep:.14f} (double-root refined)")
    print(f"E_EP = {energy_ep:.14f}")
    print(f"first coalescing sector: 2*Sz = {twice_sz}")
    print(f"final bracket: [{bracket[0]:.14f}, {bracket[1]:.14f}]")
    print(f"double-root residual ||(p,p_E)||_inf = {residual:.3e}")

    print("\n=== Puiseux coefficients for delta=gamma-gamma_EP>0 ===")
    for name, value in coefficients.items():
        print(f"{name} = {value.real:+.10f}{value.imag:+.10f}j")

    print("\n=== Jordan canonical form ===")
    jordan, jordan_blocks, nullities = jordan_form_N4(
        gamma_ep, energy_ep, twice_sz_ep=twice_sz
    )
    print(f"J shape = {jordan.shape}; nullities dim ker(H-E)={nullities[0]}, dim ker((H-E)^2)={nullities[1]}")
    for size, value, sector in jordan_blocks:
        print(f"J_{size}({value.real:+.10f}{value.imag:+.2e}j) in 2*Sz={sector}")

    print("\n=== Phase II-C: below-EP Puiseux fit in Sz=0 sector ===")
    slope, _, _ = puiseux_exponent(gamma_ep)
    print(f"Fitted exponent nu = {slope:.6f} (expected 0.5 for a second-order EP)")