"""core.py -- pure-state parametrizations for biseparable 4-qubit states.

state_1_3(params, source): pure state, product across `source` | (other 3 qubits).
state_2_2(params, pair1, pair2): pure state, product across pair1 | pair2 (each a
2-qubit tuple of indices).

Qubit order throughout this project is (A,B,C,D) = (0,1,2,3), and a 16-dim state
vector is indexed linearly as idx = i0*8 + i1*4 + i2*2 + i3.
"""
import numpy as np


def random_pure_state(dim, rng):
    v = rng.normal(size=dim) + 1j * rng.normal(size=dim)
    return v / np.linalg.norm(v)


def state_1_3(params, source):
    """Pure state product across `source` (single qubit) | (other 3 qubits).
    params: 18 reals = 2 (single-qubit Bloch angles theta,phi) + 16 (3-qubit target
    complex amplitudes, given as 8 real + 8 imaginary parts)."""
    theta, phi = params[0], params[1]
    q_src = np.array([np.cos(theta / 2), np.exp(1j * phi) * np.sin(theta / 2)], dtype=complex)
    tgt_re = params[2:2 + 8]
    tgt_im = params[2 + 8:2 + 16]
    q_tgt = tgt_re + 1j * tgt_im
    q_tgt = q_tgt / np.linalg.norm(q_tgt)
    others = [k for k in range(4) if k != source]
    full = np.zeros(16, dtype=complex)
    for s in range(2):
        for t_idx in range(8):
            bits = [(t_idx >> 2) & 1, (t_idx >> 1) & 1, t_idx & 1]
            idx = [0, 0, 0, 0]
            idx[source] = s
            for oi, o in enumerate(others):
                idx[o] = bits[oi]
            lin = idx[0] * 8 + idx[1] * 4 + idx[2] * 2 + idx[3]
            full[lin] = q_src[s] * q_tgt[t_idx]
    return full


def state_2_2(params, pair1, pair2):
    """Pure state product across pair1 | pair2 (each a 2-qubit index tuple).
    params: 16 reals = 8 (pair1 complex amplitudes) + 8 (pair2 complex amplitudes)."""
    p1 = params[0:4] + 1j * params[4:8]
    p1 = p1 / np.linalg.norm(p1)
    p2 = params[8:12] + 1j * params[12:16]
    p2 = p2 / np.linalg.norm(p2)
    full = np.zeros(16, dtype=complex)
    for a in range(2):
        for b in range(2):
            for c in range(2):
                for d in range(2):
                    idx = [0, 0, 0, 0]
                    idx[pair1[0]] = a
                    idx[pair1[1]] = b
                    idx[pair2[0]] = c
                    idx[pair2[1]] = d
                    lin = idx[0] * 8 + idx[1] * 4 + idx[2] * 2 + idx[3]
                    full[lin] = p1[a * 2 + b] * p2[c * 2 + d]
    return full