
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  PKTron Quantum GAN (QGAN)
#
#  What is a GAN?
#    A Generative Adversarial Network has two parts:
#      Generator  → tries to CREATE fake data that looks real
#      Discriminator → tries to TELL APART real vs fake data
#    They compete against each other until the generator
#    produces data indistinguishable from real data.
#
#  What is a QGAN?
#    Both the Generator and Discriminator are QUANTUM CIRCUITS
#    with trainable rotation angles (parameters).
#    The circuits are trained using gradient descent —
#    specifically the Parameter Shift Rule (a quantum gradient).
#
#  This module gives you:
#    QGAN          — main trainable class
#    QGANResult    — stores training history & generated samples
#    plot_qgan     — visualize training and generated data
#
#  Usage:
#      qgan = QGAN(n_qubits=3, n_layers=2)
#      result = qgan.train(real_data, epochs=50)
#      samples = qgan.generate(n_samples=100)
#      plot_qgan(result)
# ─────────────────────────────────────────────────────────────────


# ── Low-level gate math (no PKTron import needed inside circuits) ─

def _rx(t):
    c,s = np.cos(t/2), np.sin(t/2)
    return np.array([[c,-1j*s],[-1j*s,c]], dtype=complex)

def _ry(t):
    c,s = np.cos(t/2), np.sin(t/2)
    return np.array([[c,-s],[s,c]], dtype=complex)

def _rz(t):
    return np.array([[np.exp(-1j*t/2),0],
                     [0,np.exp(1j*t/2)]], dtype=complex)

def _H():
    return np.array([[1,1],[1,-1]], dtype=complex) / np.sqrt(2)

def _CNOT():
    return np.array([[1,0,0,0],[0,1,0,0],
                     [0,0,0,1],[0,0,1,0]], dtype=complex)

def _apply_1q(sv, mat, q, n):
    psi   = sv.reshape([2]*n)
    psi   = np.tensordot(mat, psi, axes=[[1],[q]])
    order = list(range(1,q+1)) + [0] + list(range(q+1,n))
    return np.transpose(psi, order).reshape(-1)

def _apply_2q(sv, mat, q0, q1, n):
    psi = sv.reshape([2]*n)
    G   = mat.reshape(2,2,2,2)
    ii  = list(range(n)); oi = list(range(n))
    ni,nj = n, n+1
    oi[q0]=ni; oi[q1]=nj
    return np.einsum(G,[ni,nj,q0,q1], psi,ii, oi).reshape(-1)


# ── Parameterised quantum circuit ─────────────────────────────────

def _run_circuit(n_qubits: int, params: np.ndarray,
                 input_vec: np.ndarray = None) -> np.ndarray:
    """
    Run a variational quantum circuit and return the statevector.

    Architecture (one layer):
      1. H on every qubit         (creates superposition)
      2. Ry(params[i])            (trainable rotation)
      3. Rz(params[n+i])          (trainable phase)
      4. CNOT chain               (entanglement)

    params shape: (n_layers * 2 * n_qubits,)
    """
    n  = n_qubits
    sv = np.zeros(2**n, dtype=complex)
    sv[0] = 1.0

    # Optional: encode input data as Rx rotations
    if input_vec is not None:
        for q in range(min(n, len(input_vec))):
            sv = _apply_1q(sv, _rx(float(input_vec[q])), q, n)

    # Count layers from params size
    n_params_per_layer = 2 * n
    n_layers = max(1, len(params) // n_params_per_layer)

    for layer in range(n_layers):
        offset = layer * n_params_per_layer

        # H layer
        for q in range(n):
            sv = _apply_1q(sv, _H(), q, n)

        # Ry rotations
        for q in range(n):
            idx = offset + q
            if idx < len(params):
                sv = _apply_1q(sv, _ry(params[idx]), q, n)

        # Rz rotations
        for q in range(n):
            idx = offset + n + q
            if idx < len(params):
                sv = _apply_1q(sv, _rz(params[idx]), q, n)

        # CNOT entanglement chain: 0→1, 1→2, ..., (n-2)→(n-1)
        for q in range(n - 1):
            sv = _apply_2q(sv, _CNOT(), q, q+1, n)

    return sv


def _measure_expectation(sv: np.ndarray) -> float:
    """
    Measure expectation value of Z on qubit 0.
    <Z_0> = sum_i sign(i) * |amplitude_i|^2
    where sign(i) = +1 if bit 0 of i is 0, -1 if bit 0 is 1.
    Range: [-1, +1]
    """
    probs = np.abs(sv)**2
    signs = np.array([1 if (i & 1) == 0 else -1
                      for i in range(len(probs))])
    return float(np.dot(signs, probs))


def _sample_from_sv(sv: np.ndarray, n_samples: int,
                    n_qubits: int) -> np.ndarray:
    """
    Sample bitstrings from statevector and convert to floats in [-1,1].
    """
    probs   = np.abs(sv)**2
    probs  /= probs.sum()
    indices = np.random.choice(len(probs), size=n_samples, p=probs)
    # Convert each bitstring index to a float in [-1, 1]
    samples = np.array([
        (2 * int(format(idx, f'0{n_qubits}b')[0]) - 1)
        for idx in indices
    ], dtype=float)
    return samples


# ── Parameter Shift Rule (quantum gradient) ───────────────────────

def _param_shift_grad(params: np.ndarray,
                      loss_fn,
                      shift: float = np.pi/2) -> np.ndarray:
    """
    Compute gradient using the Parameter Shift Rule.
    grad[i] = (loss(params[i] + π/2) - loss(params[i] - π/2)) / 2
    This is an EXACT gradient for quantum circuits (not finite diff).
    """
    grad = np.zeros_like(params)
    for i in range(len(params)):
        p_plus  = params.copy(); p_plus[i]  += shift
        p_minus = params.copy(); p_minus[i] -= shift
        grad[i] = (loss_fn(p_plus) - loss_fn(p_minus)) / 2.0
    return grad


# ── Result class ──────────────────────────────────────────────────

class QGANResult:
    """
    Stores all training results from a QGAN run.

    Attributes:
        g_losses      : generator loss per epoch
        d_losses      : discriminator loss per epoch
        g_params      : final generator parameters
        d_params      : final discriminator parameters
        generated     : final generated samples
        real_data     : the original real data
        n_qubits      : number of qubits used
        n_layers      : number of circuit layers
        epochs        : total training epochs
    """
    def __init__(self, g_losses, d_losses, g_params, d_params,
                 generated, real_data, n_qubits, n_layers, epochs):
        self.g_losses  = g_losses
        self.d_losses  = d_losses
        self.g_params  = g_params
        self.d_params  = d_params
        self.generated = generated
        self.real_data = real_data
        self.n_qubits  = n_qubits
        self.n_layers  = n_layers
        self.epochs    = epochs

    def summary(self):
        print("=" * 58)
        print("  PKTron QGAN — Training Summary")
        print("=" * 58)
        print(f"  Qubits          : {self.n_qubits}")
        print(f"  Layers          : {self.n_layers}")
        print(f"  Epochs trained  : {self.epochs}")
        print(f"  Generator params: {len(self.g_params)}")
        print(f"  Discrim. params : {len(self.d_params)}")
        print(f"  Final G loss    : {self.g_losses[-1]:.6f}")
        print(f"  Final D loss    : {self.d_losses[-1]:.6f}")
        print(f"  Real data mean  : {np.mean(self.real_data):.4f}")
        print(f"  Generated mean  : {np.mean(self.generated):.4f}")
        print(f"  Real data std   : {np.std(self.real_data):.4f}")
        print(f"  Generated std   : {np.std(self.generated):.4f}")
        print("=" * 58)


# ── Main QGAN class ───────────────────────────────────────────────

class QGAN:
    """
    Quantum Generative Adversarial Network.

    Both Generator and Discriminator are parameterised
    quantum circuits trained with the Parameter Shift Rule.

    Args:
        n_qubits : qubits per circuit (3-5 recommended)
        n_layers : circuit depth / layers (2-4 recommended)
        lr_g     : generator learning rate
        lr_d     : discriminator learning rate
        seed     : random seed for reproducibility

    Example:
        qgan = QGAN(n_qubits=3, n_layers=2)
        result = qgan.train(real_data, epochs=30)
        result.summary()
        samples = qgan.generate(100)
        plot_qgan(result)
    """

    def __init__(self, n_qubits: int = 3, n_layers: int = 2,
                 lr_g: float = 0.05, lr_d: float = 0.05,
                 seed: int = 42):
        self.n_qubits = n_qubits
        self.n_layers = n_layers
        self.lr_g     = lr_g
        self.lr_d     = lr_d
        self.seed     = seed

        rng = np.random.default_rng(seed)
        n_params = n_layers * 2 * n_qubits

        # Initialise parameters near zero (small random rotations)
        self.g_params = rng.uniform(-0.1, 0.1, n_params)
        self.d_params = rng.uniform(-0.1, 0.1, n_params)

        self._g_losses = []
        self._d_losses = []

        print(f"[QGAN] Initialized")
        print(f"  Qubits : {n_qubits}")
        print(f"  Layers : {n_layers}")
        print(f"  Params : {n_params} (generator) + "
              f"{n_params} (discriminator)")

    # ── Internal forward passes ───────────────────────────────────

    def _generate_sample(self, params=None) -> float:
        """Generate one sample using the generator circuit."""
        if params is None:
            params = self.g_params
        sv = _run_circuit(self.n_qubits, params)
        return _measure_expectation(sv)

    def _discriminate(self, value: float, params=None) -> float:
        """
        Discriminator: given a value, return score in [-1, 1].
        Higher = more likely to be real.
        Encode the value as an input rotation.
        """
        if params is None:
            params = self.d_params
        input_vec = np.array([value * np.pi])
        sv = _run_circuit(self.n_qubits, params, input_vec)
        return _measure_expectation(sv)

    # ── Loss functions ────────────────────────────────────────────

    def _d_loss(self, d_params: np.ndarray,
                real_batch: np.ndarray) -> float:
        """
        Discriminator loss.
        Wants: D(real) → +1, D(fake) → -1
        Loss = -mean(D(real)) + mean(D(fake))
        (minimise this → D gets better at telling real from fake)
        """
        # Score real samples (want high = close to +1)
        real_scores = np.mean([
            self._discriminate(x, d_params)
            for x in real_batch
        ])
        # Score fake samples (want low = close to -1)
        fake_samples = [self._generate_sample() for _ in real_batch]
        fake_scores  = np.mean([
            self._discriminate(x, d_params)
            for x in fake_samples
        ])
        return -real_scores + fake_scores

    def _g_loss(self, g_params: np.ndarray) -> float:
        """
        Generator loss.
        Wants: D(fake) → +1  (fool the discriminator)
        Loss = -mean(D(G(z)))
        (minimise this → G produces more convincing fakes)
        """
        fake_samples = [
            _measure_expectation(
                _run_circuit(self.n_qubits, g_params)
            )
            for _ in range(4)   # small batch per gradient step
        ]
        scores = np.mean([
            self._discriminate(x) for x in fake_samples
        ])
        return -scores

    # ── Training loop ─────────────────────────────────────────────

    def train(self, real_data: np.ndarray,
              epochs: int = 30,
              batch_size: int = 8,
              n_d_steps: int = 1,
              verbose: bool = True) -> QGANResult:
        """
        Train the QGAN.

        Args:
            real_data  : 1D numpy array of real data samples
            epochs     : number of training epochs
            batch_size : samples per gradient step
            n_d_steps  : discriminator steps per generator step
            verbose    : print progress every 5 epochs

        Returns:
            QGANResult with full training history

        Example:
            import numpy as np
            real = np.random.normal(0.5, 0.2, 200)
            result = qgan.train(real, epochs=40)
        """
        real_data = np.array(real_data, dtype=float)

        # Normalise real data to [-1, 1] for quantum circuits
        d_min, d_max = real_data.min(), real_data.max()
        if d_max - d_min > 1e-8:
            real_norm = 2 * (real_data - d_min) / (d_max - d_min) - 1
        else:
            real_norm = real_data.copy()

        self._g_losses = []
        self._d_losses = []
        rng = np.random.default_rng(self.seed)

        print(f"\n[QGAN] Training for {epochs} epochs ...")
        print(f"  Real data: {len(real_data)} samples | "
              f"mean={real_norm.mean():.3f} | "
              f"std={real_norm.std():.3f}")
        print("-" * 50)

        for epoch in range(epochs):

            # ── Discriminator update ──────────────────────────────
            for _ in range(n_d_steps):
                batch_idx = rng.choice(len(real_norm),
                                       size=min(batch_size,
                                                len(real_norm)),
                                       replace=False)
                batch = real_norm[batch_idx]

                d_grad = _param_shift_grad(
                    self.d_params,
                    lambda p: self._d_loss(p, batch)
                )
                self.d_params -= self.lr_d * d_grad

            d_loss_val = self._d_loss(self.d_params, batch)

            # ── Generator update ──────────────────────────────────
            g_grad = _param_shift_grad(
                self.g_params,
                self._g_loss
            )
            self.g_params -= self.lr_g * g_grad
            g_loss_val = self._g_loss(self.g_params)

            self._d_losses.append(float(d_loss_val))
            self._g_losses.append(float(g_loss_val))

            if verbose and (epoch % 5 == 0 or epoch == epochs-1):
                print(f"  Epoch {epoch+1:>3}/{epochs} | "
                      f"G loss: {g_loss_val:>8.5f} | "
                      f"D loss: {d_loss_val:>8.5f}")

        print("-" * 50)
        print("[QGAN] Training complete ✅")

        # Generate final samples
        final_samples = self.generate(n_samples=len(real_data))

        return QGANResult(
            g_losses  = self._g_losses,
            d_losses  = self._d_losses,
            g_params  = self.g_params.copy(),
            d_params  = self.d_params.copy(),
            generated = final_samples,
            real_data = real_data,
            n_qubits  = self.n_qubits,
            n_layers  = self.n_layers,
            epochs    = epochs,
        )

    # ── Generation ────────────────────────────────────────────────

    def generate(self, n_samples: int = 100) -> np.ndarray:
        """
        Generate samples using the trained generator.

        Returns numpy array of shape (n_samples,).

        Example:
            samples = qgan.generate(200)
            print(samples.mean(), samples.std())
        """
        samples = np.array([
            self._generate_sample() for _ in range(n_samples)
        ])
        return samples

    def save_params(self, path: str = "qgan_params.npz"):
        """Save generator and discriminator parameters to a file."""
        np.savez(path,
                 g_params=self.g_params,
                 d_params=self.d_params)
        print(f"[QGAN] Parameters saved to {path}")

    def load_params(self, path: str = "qgan_params.npz"):
        """Load previously saved parameters."""
        data = np.load(path)
        self.g_params = data["g_params"]
        self.d_params = data["d_params"]
        print(f"[QGAN] Parameters loaded from {path}")


# ── Plotting ──────────────────────────────────────────────────────

def plot_qgan(result: QGANResult):
    """
    4-panel plot showing:
      1. Generator & Discriminator loss curves
      2. Real vs Generated data distribution
      3. Generated sample evolution (first vs last)
      4. Parameter norms over training (approximated)

    Args:
        result: QGANResult from qgan.train()
    """
    import matplotlib.pyplot as plt

    fig, axes = plt.subplots(2, 2, figsize=(13, 9))

    epochs = np.arange(1, len(result.g_losses) + 1)

    # ── Panel 1: Loss curves ──────────────────────────────────────
    ax = axes[0, 0]
    ax.plot(epochs, result.g_losses,
            label='Generator Loss', color='steelblue', linewidth=1.8)
    ax.plot(epochs, result.d_losses,
            label='Discriminator Loss', color='tomato',
            linewidth=1.8, linestyle='--')
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Loss")
    ax.set_title("Training Loss Curves")
    ax.legend()
    ax.grid(alpha=0.3)

    # ── Panel 2: Real vs Generated distribution ───────────────────
    ax2 = axes[0, 1]
    ax2.hist(result.real_data, bins=25, alpha=0.6,
             color='steelblue', label='Real Data', density=True)
    ax2.hist(result.generated, bins=25, alpha=0.6,
             color='tomato', label='Generated Data', density=True)
    ax2.set_xlabel("Value")
    ax2.set_ylabel("Density")
    ax2.set_title("Real vs Generated Distribution")
    ax2.legend()
    ax2.grid(alpha=0.3)

    # ── Panel 3: Generator loss smoothed ─────────────────────────
    ax3 = axes[1, 0]
    window = max(1, len(result.g_losses) // 10)
    smooth = np.convolve(result.g_losses,
                         np.ones(window)/window, mode='valid')
    ax3.plot(smooth, color='steelblue', linewidth=2,
             label=f'G Loss (smoothed, w={window})')
    ax3.fill_between(range(len(smooth)), smooth,
                     alpha=0.2, color='steelblue')
    ax3.set_xlabel("Epoch")
    ax3.set_ylabel("Generator Loss")
    ax3.set_title("Generator Convergence")
    ax3.legend()
    ax3.grid(alpha=0.3)

    # ── Panel 4: Generated sample statistics ─────────────────────
    ax4 = axes[1, 1]
    stats = {
        'Real\nMean':  np.mean(result.real_data),
        'Gen\nMean':   np.mean(result.generated),
        'Real\nStd':   np.std(result.real_data),
        'Gen\nStd':    np.std(result.generated),
        'Real\nMin':   np.min(result.real_data),
        'Gen\nMin':    np.min(result.generated),
        'Real\nMax':   np.max(result.real_data),
        'Gen\nMax':    np.max(result.generated),
    }
    colors = ['steelblue','tomato'] * 4
    bars = ax4.bar(stats.keys(), stats.values(),
                   color=colors, edgecolor='navy',
                   linewidth=0.5)
    ax4.set_title("Real vs Generated — Statistics")
    ax4.set_ylabel("Value")
    ax4.grid(axis='y', alpha=0.3)
    ax4.tick_params(axis='x', labelsize=8)

    plt.suptitle(
        f"PKTron QGAN — {result.n_qubits} qubits | "
        f"{result.n_layers} layers | "
        f"{result.epochs} epochs",
        fontsize=12, fontweight='bold'
    )
    plt.tight_layout()
    plt.savefig('pktron_qgan.png', dpi=120, bbox_inches='tight')
    plt.show()
    print("[QGAN] Plot saved as pktron_qgan.png")
