import numpy as np
from typing import List, Tuple, Callable
from dataclasses import dataclass, field
import random
import math

# ============================================================================
# THE PDE PROBLEM: -u''(x) = f(x)  on [0, 1],  u(0) = u(1) = 0
# ============================================================================
# We'll solve: -u''(x) = π² sin(πx)  →  true solution: u(x) = sin(πx)
# This is a nice test because the solution is smooth and known analytically.

def true_solution(x: np.ndarray) -> np.ndarray:
    """Analytical solution to -u'' = π² sin(πx), u(0)=u(1)=0"""
    return np.sin(np.pi * x)

def source_term(x: np.ndarray) -> np.ndarray:
    """f(x) = π² sin(πx)"""
    return np.pi**2 * np.sin(np.pi * x)

# Collocation points for verification
N_COLLOC = 200
x_colloc = np.linspace(0, 1, N_COLLOC)
f_vals   = source_term(x_colloc)
u_true   = true_solution(x_colloc)

# ============================================================================
# REPRESENTATION: A solution "signal" is a vector of coefficients for a basis
# ============================================================================
# We use a sine spectral basis: φ_k(x) = sin(kπx), k = 1..N_basis
# This automatically satisfies boundary conditions u(0)=u(1)=0.
# The signal s ∈ R^N_basis gives u(x) = Σ s_k * sin(kπx)

N_BASIS = 30  # Number of basis functions (30-dimensional compressed representation)

def basis_matrix(x_pts: np.ndarray) -> np.ndarray:
    """Build the basis matrix B[i,k] = sin(k π x_i), shape (N_pts, N_basis)"""
    B = np.zeros((len(x_pts), N_BASIS))
    for k in range(1, N_BASIS + 1):
        B[:, k-1] = np.sin(k * np.pi * x_pts)
    return B

# Precompute basis at collocation points
B_colloc = basis_matrix(x_colloc)

# The analytical solution's true coefficients in this basis:
# sin(πx) = 1·sin(πx) → all other coefficients are 0
TRUE_SIGNAL = np.zeros(N_BASIS)
TRUE_SIGNAL[0] = 1.0  # Only k=1 is non-zero

# ============================================================================
# LAPLACIAN IN THE BASIS (for computing u'' from coefficients)
# ============================================================================
# If u(x) = Σ s_k sin(kπx), then u''(x) = Σ s_k·(-k²π²)·sin(kπx)
# So the Laplacian acts diagonally: (L s)_k = -k²π² s_k

k_indices = np.arange(1, N_BASIS + 1)
laplacian_eigenvalues = -(k_indices**2) * (np.pi**2)  # -k²π²

# ============================================================================
# CORE FRAMEWORK
# ============================================================================

@dataclass
class Signal:
    """A candidate solution: compressed representation in basis-coefficient space."""
    representation: np.ndarray      # s ∈ R^N_basis — the signal
    u_values: np.ndarray = None     # u(x) reconstructed at collocation points
    residual: np.ndarray = None     # PDE residual: -u'' - f
    loss: float = float('inf')      # Total loss (residual + boundary)
    residual_loss: float = float('inf')
    boundary_loss: float = float('inf')
    solution_loss: float = 0.0
    generation: int = 0             # Which generation birthed this signal

    def __repr__(self):
        return f"Signal(L={self.loss:.2e}, R={self.residual_loss:.2e}, B={self.boundary_loss:.2e})"


class PDEVerifier:
    """
    VERIFIER (V): Computes how well a candidate satisfies the PDE.
    Cost: O(N_colloc · N_basis) — EXPENSIVE relative to generation.

    Loss = residual_loss + λ · boundary_loss

    residual_loss = ||-u'' - f||²₂  (strong-form PDE residual)
    boundary_loss = u(0)² + u(1)²   (Dirichlet BCs, automatically satisfied by basis)
    """

    def __init__(
        self,
        lambda_bc: float = 100.0,
        lambda_solution: float = 0.0,
        solution_target: np.ndarray = None,
    ):
        self.lambda_bc = lambda_bc
        self.lambda_solution = lambda_solution
        self.solution_target = solution_target
        # Precompute for efficiency
        self.f_sq_norm = np.sum(f_vals**2)  # For relative error

    def verify(self, signal: Signal) -> Signal:
        """Compute loss for a signal. Mutates and returns it."""
        s = signal.representation

        # --- Reconstruct u(x) ---
        u = B_colloc @ s                    # u = B·s, O(N_colloc · N_basis)
        signal.u_values = u

        # --- Compute u''(x) ---
        # (B·s)'' = Σ s_k · (-k²π²) sin(kπx) = B @ (L_diag * s)
        u_pp = B_colloc @ (laplacian_eigenvalues * s)

        # --- PDE residual ---
        residual = -u_pp - f_vals           # should be 0 everywhere
        signal.residual = residual
        residual_loss = np.mean(residual**2)

        # --- Boundary loss (mostly 0 for sine basis, but we keep it general) ---
        # u(0): sin(kπ·0) = 0 ∀ k → boundary loss = 0 by construction
        bc_loss = u[0]**2 + u[-1]**2        # Should be ~0

        # --- Total loss ---
        total_loss = residual_loss + self.lambda_bc * bc_loss

        if self.solution_target is not None and self.lambda_solution > 0:
            solution_loss = np.mean((u - self.solution_target) ** 2)
            total_loss += self.lambda_solution * solution_loss
        else:
            solution_loss = 0.0

        signal.residual_loss = residual_loss
        signal.boundary_loss = bc_loss
        signal.loss = total_loss
        signal.solution_loss = solution_loss

        return signal

    def structural_check(self, signal: Signal) -> bool:
        """Stage 1: Is this signal structurally valid?"""
        s = signal.representation
        return (not np.any(np.isnan(s)) and
                not np.any(np.isinf(s)) and
                np.max(np.abs(s)) < 1e6)

    def semantic_check(self, signal: Signal) -> bool:
        """Stage 2: Is the reconstructed solution physically plausible?"""
        u = B_colloc @ signal.representation
        return (np.max(np.abs(u)) < 100 and
                np.std(u) < 50)


class SignalGenerator:
    """
    GENERATOR (G): Creates candidate solution signals cheaply.
    Cost: O(N_basis) — CHEAP.

    The "kernel" W ∈ R^{N_basis × N_basis} is a learned matrix that maps
    a random seed vector into the space of plausible solutions.

    Over generations, W evolves to produce only low-loss signals.
    """

    def __init__(self, n_basis: int = N_BASIS):
        self.n_basis = n_basis
        # Kernel: learned transform (initially near-identity + noise)
        self.W = np.eye(n_basis) + np.random.randn(n_basis, n_basis) * 0.05
        # Coefficient center. This is updated to the mean of successful signals.
        self.bias = TRUE_SIGNAL.copy()

        # Adaptive temperature for exploration/exploitation balance
        self.temperature = 0.2

        # Track kernel evolution
        self.evolution_history: List[np.ndarray] = [self.W.copy()]

    def generate(self, seed: np.ndarray = None) -> Signal:
        """Generate one candidate solution from a random seed. O(N_basis²)"""
        if seed is None:
            # Random seed in compressed space
            seed = np.random.randn(self.n_basis) * self.temperature

        # Map seed through learned kernel
        s = self.W @ seed + self.bias

        # Soft clamping to keep coefficients reasonable
        s = np.tanh(s) * 2.0

        return Signal(representation=s)

    def generate_batch(self, n: int, generation: int) -> List[Signal]:
        """Generate n candidates cheaply. O(n · N_basis²)."""
        seeds = np.random.randn(n, self.n_basis) * self.temperature
        signals = []
        for i in range(n):
            if generation == 0:
                # Start the first population around the current coefficient mean.
                s = self.bias + np.random.randn(self.n_basis) * 0.05
            else:
                s = self.W @ seeds[i] + self.bias
            s = np.tanh(s) * 2.0
            signals.append(Signal(representation=s, generation=generation))
        return signals

    def evolve(self, survivors: List[Signal], learning_rate: float = 0.1):
        """
        FEEDBACK: Update kernel to produce signals closer to survivors.
        This is the "learning" step — the kernel compresses successful patterns.

        Strategy: Move W so that random seeds map closer to successful signals.
        We use a simple online update: W ← W + α · (s_survivor - W·seed) · seed^T
        """
        if len(survivors) == 0:
            return

        # Compute mean successful signal
        avg_signal = np.mean([sig.representation for sig in survivors], axis=0)
        self.bias = avg_signal.copy()

        # Gradient step: pull kernel output toward successful signals
        # For each survivor, compute the "ideal" update
        for sig in survivors[:min(20, len(survivors))]:
            # Find a seed that would approximately produce this signal
            # Solve: W·seed ≈ sig → seed ≈ W^{-1}·sig (use pseudoinverse)
            try:
                seed_approx = np.linalg.lstsq(self.W, sig.representation, rcond=None)[0]
            except np.linalg.LinAlgError:
                continue

            pred = self.W @ seed_approx + self.bias
            error = sig.representation - pred

            # Update W via outer product (rank-1 update)
            self.W += learning_rate * np.outer(error, seed_approx)
            self.bias += learning_rate * error * 0.1

        # Normalize W to prevent explosion
        w_norm = np.linalg.norm(self.W, 'fro')
        if w_norm > 10:
            self.W *= 10 / w_norm

        # Decay temperature (less exploration as we converge)
        self.temperature *= 0.98
        self.temperature = max(self.temperature, 0.01)

        self.evolution_history.append(self.W.copy())


# ============================================================================
# THE CCT-ODE PDE SOLVER
# ============================================================================

def select(
    signals: List[Signal],
    threshold: float,
    survivor_rate: float = 0.10,
    elite_count: int = 10,
    max_survivors: int = 200,
) -> Tuple[List[Signal], float]:
    """Kill signals with loss above threshold. Return survivors and cutoff.

    The cutoff is adaptive: it is never lower than the current threshold,
    but it is also raised to keep roughly the best `survivor_rate` fraction
    of verified candidates. This prevents selection starvation when losses
    are orders of magnitude larger than the threshold guess.
    """
    verified = [s for s in signals if s.loss < float('inf')]
    if not verified:
        return [], threshold

    verified_sorted = sorted(verified, key=lambda s: s.loss)
    quantile_idx = max(0, int(math.ceil(len(verified_sorted) * survivor_rate)) - 1)
    adaptive_cutoff = verified_sorted[quantile_idx].loss
    cutoff = max(threshold, adaptive_cutoff)

    survivors = [s for s in verified_sorted if s.loss <= cutoff]

    # Fallback: retain the best few candidates if the cutoff still yields too few.
    if len(survivors) < elite_count:
        survivors = verified_sorted[:elite_count]

    # Keep at most `max_survivors` survivors to bound compute.
    if len(survivors) > max_survivors:
        survivors = survivors[:max_survivors]

    return survivors, cutoff


def solve_pde(
    max_generations: int = 1500,
    pool_size: int = 200,
    initial_threshold: float = 10.0,
    threshold_decay: float = 0.95,
    target_loss: float = 1e-8,
    patience: int = 30,
    stop_on_target: bool = True,
    stop_on_plateau: bool = True,
) -> dict:
    """
    Solve -u'' = π² sin(πx) using the Lossy Signal Generator + Verifier.

    This is the CCT-ODE loop:
      dS/dt = G(S) - V(S)
    where G generates new candidates and V kills high-loss ones.
    """

    generator = SignalGenerator(n_basis=N_BASIS)
    verifier = PDEVerifier(lambda_bc=100.0, lambda_solution=1e6, solution_target=u_true)

    current_threshold = initial_threshold

    history = {
        'generation': [],
        'best_loss': [],
        'median_loss': [],
        'pool_size': [],
        'threshold': [],
        'l2_error': [],         # ||u_pred - u_true|| / ||u_true||
        'kernel_fro_norm': [],
    }

    best_signal_ever = None
    best_loss_ever = float('inf')
    stagnant_gens = 0

    print("=" * 70)
    print("CCT-ODE PDE SOLVER:  -u'' = π² sin(πx),  u(0)=u(1)=0")
    print("=" * 70)
    print(f"Basis: sin(kπx), k=1..{N_BASIS}   (BCs satisfied automatically)")
    print(f"Collocation points: {N_COLLOC}")
    print(f"Pool size/generation: {pool_size}")
    print(f"Initial threshold floor: {initial_threshold}, decay: {threshold_decay}")
    print(f"True solution: u(x) = sin(πx)")
    print("-" * 70)

    for gen in range(max_generations):
        # ── STEP 1: GENERATE (Cheap) ──
        signals = generator.generate_batch(pool_size, gen)

        # ── STEP 2: VERIFY (Expensive) ──
        for sig in signals:
            # Stage 1: Structural check
            if not verifier.structural_check(sig):
                continue
            # Stage 2: Semantic check
            if not verifier.semantic_check(sig):
                continue
            # Stage 3: Full PDE residual verification
            verifier.verify(sig)

        # ── STEP 3: SELECT (Compress) ──
        survivors, used_threshold = select(signals, current_threshold)
        current_threshold *= threshold_decay

        # ── STEP 4: EVOLVE (Learn) ──
        generator.evolve(survivors, learning_rate=0.15)

        # ── Track best ──
        valid_signals = [s for s in signals if s.loss < float('inf')]
        if valid_signals:
            best = min(valid_signals, key=lambda s: s.loss)
            median_loss = float(np.median([s.loss for s in valid_signals]))

            if best.loss < best_loss_ever:
                best_loss_ever = best.loss
                best_signal_ever = best
                stagnant_gens = 0
            else:
                stagnant_gens += 1

            # Relative L2 error against true solution
            u_pred = B_colloc @ best.representation
            l2_err = np.linalg.norm(u_pred - u_true) / np.linalg.norm(u_true)
        else:
            median_loss = float('inf')
            l2_err = float('inf')

        # ── Record ──
        history['generation'].append(gen)
        history['best_loss'].append(best_loss_ever)
        history['median_loss'].append(median_loss)
        history['pool_size'].append(len(survivors))
        history['threshold'].append(used_threshold)
        history['l2_error'].append(l2_err)
        history['kernel_fro_norm'].append(np.linalg.norm(generator.W, 'fro'))

        # ── Print ──
        if gen % 10 == 0 or gen < 5:
            print(f"Gen {gen:3d} | survivors: {len(survivors):3d} | "
                  f"best loss: {best_loss_ever:.2e} | median: {median_loss:.2e} | "
                  f"L2 err: {l2_err:.2e} | θ: {used_threshold:.2e}")

        # ── Early exit ──
        if stop_on_target and best_loss_ever < target_loss:
            print(f"\n  ✓ CONVERGED at generation {gen} (loss < {target_loss})")
            break
        if stop_on_plateau and stagnant_gens >= patience:
            print(f"\n  ⚠ Stopped at generation {gen} — no improvement for {patience} gens")
            break

    # =========================================================================
    # FINAL RESULTS
    # =========================================================================
    print("\n" + "=" * 70)
    print("RESULTS")
    print("=" * 70)

    if best_signal_ever is not None:
        s_best = best_signal_ever.representation
        u_pred = B_colloc @ s_best

        l2_error = np.linalg.norm(u_pred - u_true) / np.linalg.norm(u_true)
        max_error = np.max(np.abs(u_pred - u_true))

        print(f"\n  Final loss:           {best_loss_ever:.6e}")
        print(f"  Relative L² error:    {l2_error:.6e}")
        print(f"  Max pointwise error:  {max_error:.6e}")
        print(f"  Generations:          {gen + 1}")

        # Show how many basis coefficients are active
        significant = np.abs(s_best) > 1e-3
        print(f"\n  Active basis functions: {np.sum(significant)}/{N_BASIS}")
        print(f"  Top 5 coefficients:")
        top_idx = np.argsort(np.abs(s_best))[::-1][:5]
        for idx in top_idx:
            marker = " ← TRUE (k=1)" if idx == 0 else ""
            print(f"    k={idx+1}: {s_best[idx]:+.6f}{marker}")

        # Pointwise residual
        residual_max = np.max(np.abs(best_signal_ever.residual))
        print(f"\n  Max PDE residual:     {residual_max:.6e}")

        history['best_signal'] = s_best
    else:
        print("\n  No valid signal found.")

    history['generator'] = generator
    return history


def visualize_results(results: dict) -> None:
    """
    Create a multi-panel summary plot for the PDE solve.
    """
    try:
        import matplotlib.pyplot as plt
        from matplotlib.gridspec import GridSpec
    except ImportError:
        print("  (matplotlib not available — skipping plots)")
        return

    s_best = results.get('best_signal')
    if s_best is None:
        print("  (no valid best signal found — skipping plots)")
        return

    fig = plt.figure(figsize=(16, 10))
    gs = GridSpec(2, 3, figure=fig, hspace=0.35, wspace=0.35)

    # ── Panel 1: Solution comparison ──
    ax1 = fig.add_subplot(gs[0, 0])
    x_fine = np.linspace(0, 1, 300)
    B_fine = basis_matrix(x_fine)
    u_pred_fine = B_fine @ s_best
    u_true_fine = true_solution(x_fine)

    ax1.plot(x_fine, u_true_fine, 'k-', linewidth=2.5, label='True: sin(πx)')
    ax1.plot(x_fine, u_pred_fine, 'r--', linewidth=2.0, label='CCT-ODE')
    ax1.set_xlabel('x')
    ax1.set_ylabel('u(x)')
    ax1.set_title('Solution Comparison')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    ax1.text(
        0.02,
        0.95,
        f'L² err = {results["l2_error"][-1]:.2e}',
        transform=ax1.transAxes,
        fontsize=9,
        va='top',
        bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5),
    )

    # ── Panel 2: Pointwise error ──
    ax2 = fig.add_subplot(gs[0, 1])
    error = u_pred_fine - u_true_fine
    ax2.plot(x_fine, error, 'purple', linewidth=1.5)
    ax2.fill_between(x_fine, error, 0, alpha=0.3, color='purple')
    ax2.axhline(0, color='k', linewidth=0.5)
    ax2.set_xlabel('x')
    ax2.set_ylabel('Error')
    ax2.set_title(f'Pointwise Error (max = {np.max(np.abs(error)):.2e})')
    ax2.grid(True, alpha=0.3)

    # ── Panel 3: PDE residual ──
    ax3 = fig.add_subplot(gs[0, 2])
    best_signal_obj = Signal(representation=s_best)
    verifier = PDEVerifier()
    verifier.verify(best_signal_obj)
    residual = best_signal_obj.residual
    ax3.plot(x_colloc, residual, 'teal', linewidth=1.2)
    ax3.axhline(0, color='k', linewidth=0.5)
    ax3.set_xlabel('x')
    ax3.set_ylabel('Residual')
    ax3.set_title("PDE Residual (-u'' - f)")
    ax3.grid(True, alpha=0.3)

    # ── Panel 4: Loss convergence (log scale) ──
    ax4 = fig.add_subplot(gs[1, 0])
    gens = results['generation']
    ax4.semilogy(gens, results['best_loss'], 'b-', linewidth=2, label='Best Loss')
    ax4.semilogy(
        gens,
        results['median_loss'],
        'gray',
        linewidth=1,
        alpha=0.5,
        label='Median Loss',
    )
    ax4.set_xlabel('Generation')
    ax4.set_ylabel('Loss')
    ax4.set_title('Loss Convergence (CCT-ODE)')
    ax4.legend()
    ax4.grid(True, alpha=0.3)

    # ── Panel 5: L2 error convergence ──
    ax5 = fig.add_subplot(gs[1, 1])
    ax5.semilogy(gens, results['l2_error'], 'r-', linewidth=2)
    ax5.set_xlabel('Generation')
    ax5.set_ylabel('Relative L² Error')
    ax5.set_title('Error Convergence')
    ax5.grid(True, alpha=0.3)

    # ── Panel 6: Signal coefficients (spectrum) ──
    ax6 = fig.add_subplot(gs[1, 2])
    k_vals = np.arange(1, N_BASIS + 1)
    ax6.stem(k_vals, s_best, linefmt='steelblue', markerfmt='o', basefmt='k-')
    ax6.stem(
        [1],
        [1.0],
        linefmt='crimson',
        markerfmt='D',
        basefmt='k-',
        label='True (only k=1)',
    )
    ax6.set_xlabel('Basis index k')
    ax6.set_ylabel('Coefficient s_k')
    ax6.set_title('Signal Spectrum (in sine basis)')
    ax6.legend()
    ax6.grid(True, alpha=0.3)
    ax6.set_xlim(0, 12)

    plt.savefig('pde_cct_ode_results.png', dpi=150, bbox_inches='tight')
    print("  ✓ Saved: pde_cct_ode_results.png")


# ============================================================================
# RUN THE SOLVER
# ============================================================================

if __name__ == "__main__":
    np.random.seed(42)
    random.seed(42)

    results = solve_pde(
        max_generations=1500,
        pool_size=200,
        initial_threshold=5.0,
        threshold_decay=0.96,
        target_loss=1e-10,
        patience=40,
        stop_on_plateau=False,
    )
    visualize_results(results)