import numpy as np
from scipy.io.wavfile import read
from typing import Tuple, Optional, Callable
from dataclasses import dataclass
import logging
import time

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


# =============================================================================
# CORE DATA STRUCTURES
# =============================================================================

@dataclass
class TransformConfig:
    """Configuration for the Adversarial Reference Transform."""
    input_shape: Tuple[int, int, int]          # (H, W, C)
    reference_channels: int = 64               # Number of reference channels
    reference_grid_size: Tuple[int, int] = (3, 3)  # Grid of reference positions
    contemplation_depth: int = 5               # Number of adversarial iterations
    learning_rate: float = 0.001
    momentum: float = 0.9
    regularization: float = 1e-4
    progress_log_interval: int = 196           # Log every N spatial positions


@dataclass
class ReferenceBank:
    """Stores learned reference points and their activations."""
    positions: np.ndarray                      # (N_ref, 2) normalized coordinates
    features: np.ndarray                       # (N_ref, C) learned features
    attention_weights: np.ndarray              # (N_ref,) importance scores


class AdversarialReferenceTransform:
    """
    Non-convolutional reference transform inspired by Conv2D.
    
    Instead of sliding kernels, this transform:
    1. Maintains a bank of reference points (learned positions)
    2. Uses adversarial contemplation to select optimal references
    3. Transforms input through reference-based interpolation
    
    Key Insight: Conv2D's power comes from local receptive fields.
    We achieve similar behavior by:
    - Learning what positions matter (generator)
    - Judging if transform is "conv-like" (discriminator)
    - Iterative refinement through adversarial training
    """
    
    def __init__(
        self,
        input_shape: Tuple[int, int, int],
        reference_channels: int = 64,
        reference_grid_size: Tuple[int, int] = (3, 3),
        contemplation_depth: int = 5
    ):
        self.config = TransformConfig(
            input_shape=input_shape,
            reference_channels=reference_channels,
            reference_grid_size=reference_grid_size,
            contemplation_depth=contemplation_depth
        )
        
        self._initialize_components()
        
    def _initialize_components(self):
        """Initialize all learnable components."""
        H, W, C = self.config.input_shape
        
        # ----- Reference Bank (learnable positions + features) -----
        # Grid of initial reference positions
        grid_y, grid_x = np.meshgrid(
            np.linspace(0, 1, self.config.reference_grid_size[0]),
            np.linspace(0, 1, self.config.reference_grid_size[1]),
            indexing='ij'
        )
        initial_positions = np.stack([grid_x.ravel(), grid_y.ravel()], axis=1)
        
        # Add extra scattered references
        n_extra = self.config.reference_channels - len(initial_positions)
        if n_extra > 0:
            extra_positions = np.random.rand(n_extra, 2)
            initial_positions = np.vstack([initial_positions, extra_positions])
        
        self.reference_bank = ReferenceBank(
            positions=initial_positions.astype(np.float32),
            features=np.random.randn(len(initial_positions), self.config.reference_channels) * 0.01,
            attention_weights=np.ones(len(initial_positions), dtype=np.float32)
        )
        
        # ----- Transformation MLP (Generator) -----
        # Learns to transform input based on reference comparisons
        self.transform_mlp = MLPClassifier(
            input_size=C * 3,  # Input + 2 reference features
            hidden_size=128,
            output_size=C,
            learning_rate=np.array([0.01, 0.01, 0.01, 0.01])
        )
        
        # ----- Discriminator (Judges "conv-likeness") -----
        self.discriminator = Discriminator(input_channels=C * 2)
        
        # ----- Contemplation state -----
        self.contemplation_history = []
        
    def _compute_reference_features(self, X: np.ndarray, positions: np.ndarray) -> np.ndarray:
        """
        Extract features at reference positions via bilinear interpolation.
        
        This replaces conv2d's local receptive field with reference-based sampling.
        
        Args:
            X: Input tensor (B, H, W, C)
            positions: Reference positions (N_ref, 2) in [0, 1]
            
        Returns:
            Features at references (B, N_ref, C)
        """
        B, H, W, C = X.shape
        N_ref = positions.shape[0]
        
        # Convert normalized positions to pixel coordinates
        px = positions[:, 0] * (W - 1)  # (N_ref,)
        py = positions[:, 1] * (H - 1)  # (N_ref,)
        
        # Bilinear interpolation for each reference point
        x0 = np.floor(px).astype(int)
        y0 = np.floor(py).astype(int)
        x1 = np.clip(x0 + 1, 0, W - 1)
        y1 = np.clip(y0 + 1, 0, H - 1)
        
        wx = px - x0  # (N_ref,)
        wy = py - y0  # (N_ref,)
        
        features = np.zeros((B, N_ref, C), dtype=np.float32)
        
        for b in range(B):
            for n in range(N_ref):
                # Bilinear interpolation
                v00 = X[b, y0[n], x0[n]]
                v01 = X[b, y1[n], x0[n]]
                v10 = X[b, y0[n], x1[n]]
                v11 = X[b, y1[n], x1[n]]
                
                features[b, n] = (
                    v00 * (1 - wx[n]) * (1 - wy[n]) +
                    v01 * (1 - wx[n]) * wy[n] +
                    v10 * wx[n] * (1 - wy[n]) +
                    v11 * wx[n] * wy[n]
                )
        
        return features
    
    def _compute_spatial_context(self, X: np.ndarray) -> np.ndarray:
        """
        Compute spatial context features (gradient, laplacian approximations).
        
        Mimics Conv2D's edge detection capability without convolutions.
        """
        B, H, W, C = X.shape
        
        # Horizontal gradient (Sobel-like)
        grad_x = np.zeros_like(X)
        grad_x[:, :, 1:-1] = X[:, :, 2:] - X[:, :, :-2]
        grad_x[:, :, 0] = X[:, :, 1] - X[:, :, 0]
        grad_x[:, :, -1] = X[:, :, -1] - X[:, :, -2]
        
        # Vertical gradient
        grad_y = np.zeros_like(X)
        grad_y[:, 1:-1, :] = X[:, 2:] - X[:, :-2]
        grad_y[:, 0, :] = X[:, 1] - X[:, 0]
        grad_y[:, -1, :] = X[:, -1] - X[:, -2]
        
        # Laplacian approximation
        laplacian = np.zeros_like(X)
        laplacian[:, 1:-1, 1:-1] = (
            X[:, 2:, 1:-1] + X[:, :-2, 1:-1] +
            X[:, 1:-1, 2:] + X[:, 1:-1, :-2] -
            4 * X[:, 1:-1, 1:-1]
        )
        
        return np.concatenate([X, grad_x, grad_y, laplacian], axis=-1)
    
    def _adversarial_contemplation(
        self,
        X: np.ndarray,
        reference_features: np.ndarray,
        epoch: int,
        batch_index: Optional[int] = None,
        total_batches: Optional[int] = None
    ) -> np.ndarray:
        """
        Adversarial contemplation: Generator refines references, Discriminator judges.
        
        This iterative process forces the reference transform to capture local
        receptive field behavior without explicit convolutions.
        
        Args:
            X: Input (B, H, W, C)
            reference_features: Features at current references (B, N_ref, C)
            epoch: Current training epoch
            
        Returns:
            Refined reference features (B, H, W, C)
        """
        B, H, W, C = X.shape
        spatial_context = self._compute_spatial_context(X)
        
        for contemplation_step in range(self.config.contemplation_depth):
            step_start = time.perf_counter()
            # ----- Generator Step: Refine references to capture local patterns -----
            # Flatten spatial dimensions for MLP processing
            X_flat = X.reshape(B, H * W, C)
            
            # Randomly sample references to process
            n_ref = min(reference_features.shape[1], 16)
            ref_idx = np.random.choice(reference_features.shape[1], n_ref, replace=False)
            sampled_refs = reference_features[:, ref_idx, :]  # (B, n_ref, C)
            
            # For each spatial position, compute transform based on nearby references
            # This mimics conv2d's local processing
            output = np.zeros((B, H * W, C), dtype=np.float32)
            
            for pos_idx in range(H * W):
                pos_y = pos_idx // W
                pos_x = pos_idx % W
                
                # Find references close to this position (local receptive field)
                pos_norm = np.array([[pos_x / W, pos_y / H]], dtype=np.float32)
                distances = np.linalg.norm(
                    self.reference_bank.positions[:, :2] - pos_norm,
                    axis=1
                )
                local_refs = np.argsort(distances)[:n_ref]
                
                # Create input: [pixel, reference_features, spatial_context]
                pixel = X_flat[:, pos_idx:pos_idx+1, :]  # (B, 1, C)
                ref_features = reference_features[:, local_refs, :].mean(axis=1, keepdims=True)  # (B, 1, C)
                context = spatial_context[:, pos_y, pos_x:pos_x+1, :C]  # (B, 1, C)
                
                # Combine features
                combined = np.concatenate([pixel, ref_features, context], axis=-1)  # (B, 1, 3C)
                
                # Transform via MLP
                for b in range(B):
                    combined_b = combined[b, 0]  # (3C,)
                    transformed = self.transform_mlp.forward(combined_b.reshape(1, -1))
                    output[b, pos_idx] = transformed

                if (
                    (pos_idx + 1) % self.config.progress_log_interval == 0
                    or pos_idx + 1 == H * W
                ):
                    elapsed = time.perf_counter() - step_start
                    positions_done = pos_idx + 1
                    positions_per_second = positions_done / max(elapsed, 1e-9)
                    if batch_index is not None and total_batches is not None:
                        batch_label = f"{batch_index + 1}/{total_batches}"
                    else:
                        batch_label = "?"
                    logger.info(
                        "ART epoch %d batch %s contemplation %d/%d: %d/%d positions, %.1f pos/s",
                        epoch + 1,
                        batch_label,
                        contemplation_step + 1,
                        self.config.contemplation_depth,
                        positions_done,
                        H * W,
                        positions_per_second,
                    )
            
            # ----- Discriminator Step: Judge "conv-likeness" -----
            # Create ground truth conv-like features (for comparison)
            conv_like_features = self._mimic_conv2d(X)
            
            # Discriminator learns to distinguish real conv outputs from ref transforms
            d_loss, d_accuracy = self.discriminator.train_step(
                real=conv_like_features,
                fake=output,
                learning_rate=0.001
            )
            
            # Generator tries to fool discriminator
            if d_accuracy > 0.7:  # If discriminator is too good, improve generator
                self._improve_generator(output, conv_like_features)
            
            # ----- Refine reference positions based on contemplation -----
            self._update_reference_positions(X, output)
            
            logger.debug(
                f"Contemplation {contemplation_step}: D_loss={d_loss:.4f}, D_acc={d_accuracy:.3f}"
            )
        
        return output.reshape(B, H, W, C)
    
    def _mimic_conv2d(self, X: np.ndarray) -> np.ndarray:
        """
        Create a pseudo-conv2d output for discriminator comparison.
        
        Uses reference-based aggregation to mimic convolution behavior.
        """
        B, H, W, C = X.shape
        
        # Simple 3x3 reference aggregation (mimics conv kernel)
        kernel_size = 3
        pad = kernel_size // 2
        
        # Pad input
        X_padded = np.pad(X, ((0, 0), (pad, pad), (pad, pad), (0, 0)), mode='reflect')
        
        # For each position, average neighbors (like conv with all-ones kernel)
        output = np.zeros((B, H, W, C), dtype=np.float32)
        for dy in range(kernel_size):
            for dx in range(kernel_size):
                output += X_padded[:, dy:dy+H, dx:dx+W, :]
        
        output /= kernel_size * kernel_size
        return output
    
    def _update_reference_positions(self, X: np.ndarray, transformed: np.ndarray):
        """
        Update reference positions based on reconstruction quality.
        
        References that help reconstruct the transformed output move closer to
        important positions; others move away (adversarial positioning).
        """
        B, H, W, C = X.shape
        if transformed.ndim == 3:
            transformed = transformed.reshape(B, H, W, C)
        elif transformed.ndim != 4:
            raise ValueError(
                f"Expected transformed to have 3 or 4 dimensions, got shape {transformed.shape}"
            )
        
        # Compute reconstruction error at each reference
        errors = np.zeros(len(self.reference_bank.positions))
        
        for n, pos in enumerate(self.reference_bank.positions):
            px, py = int(pos[0] * (W - 1)), int(pos[1] * (H - 1))
            
            # Compute how much this reference contributes to final output
            # Higher contribution = more important position
            contribution = np.mean(np.abs(transformed[:, py, px, :]))
            errors[n] = contribution
        
        # Move references toward low-error regions (where they help most)
        # and away from high-error regions (adversarial adjustment)
        gradient = errors - errors.mean()
        position_updates = gradient * 0.01
        
        self.reference_bank.positions += position_updates[:, np.newaxis] * np.array([[1, 1]])
        self.reference_bank.positions = np.clip(self.reference_bank.positions, 0, 1)
        
        # Update attention weights based on contribution
        self.reference_bank.attention_weights *= (1 + errors * 0.1)
        self.reference_bank.attention_weights /= self.reference_bank.attention_weights.sum()
    
    def _improve_generator(self, fake: np.ndarray, real: np.ndarray):
        """
        Improve generator when discriminator is too strong.
        """
        # Compute difference and backprop
        diff = real - fake
        mean_diff = np.mean(diff, axis=tuple(range(diff.ndim - 1)), keepdims=False)
        target = np.asarray(mean_diff, dtype=np.float32).reshape(1, -1)
        target = np.clip(target, 1e-6, None)
        target /= target.sum(axis=1, keepdims=True)

        # Align the ART transform MLP using an input shape it actually expects.
        for _ in range(3):
            X_sample = np.random.rand(1, self.transform_mlp.W1.shape[0]).astype(np.float32)
            self.transform_mlp.update(X_sample, target)
    
    def forward(
        self,
        X: np.ndarray,
        epoch: int = 0,
        batch_index: Optional[int] = None,
        total_batches: Optional[int] = None
    ) -> np.ndarray:
        """
        Forward pass through the adversarial reference transform.
        
        Args:
            X: Input tensor (B, H, W, C)
            
        Returns:
            Transformed output (B, H, W, C)
        """
        B, H, W, C = X.shape
        
        # Step 1: Extract features at current reference positions
        reference_features = self._compute_reference_features(
            X, 
            self.reference_bank.positions
        )
        
        # Step 2: Adversarial contemplation to refine transform
        transformed = self._adversarial_contemplation(
            X,
            reference_features,
            epoch=epoch,
            batch_index=batch_index,
            total_batches=total_batches
        )
        
        # Step 3: Apply learned attention weighting
        attention = self.reference_bank.attention_weights
        attention_reshaped = attention.reshape(
            1, 1, 1, len(attention)
        ).repeat(C, axis=-1)[:, :, :, :transformed.shape[-1]]
        
        # Weighted combination of reference-based features
        output = transformed * np.mean(attention)
        
        return output
    
    def fit(self, X_train: np.ndarray, epochs: int = 10, batch_size: int = 32):
        """
        Fit the transform using adversarial contemplation.
        
        Args:
            X_train: Training data (N, H, W, C)
            epochs: Number of training epochs
            batch_size: Batch size for contemplation
        """
        logger.info(f"Starting adversarial contemplation training for {epochs} epochs")
        
        N = X_train.shape[0]
        total_batches = max((N + batch_size - 1) // batch_size, 1)
        
        for epoch in range(epochs):
            indices = np.random.permutation(N)
            epoch_loss = 0
            batch_count = 0
            epoch_start = time.perf_counter()
            
            for batch_start in range(0, N, batch_size):
                batch_index = batch_start // batch_size
                batch_idx = indices[batch_start:batch_start + batch_size]
                X_batch = X_train[batch_idx]
                batch_start_time = time.perf_counter()
                
                # Forward pass with contemplation
                output = self.forward(
                    X_batch,
                    epoch=epoch,
                    batch_index=batch_index,
                    total_batches=total_batches
                )
                
                # Compute reconstruction loss
                recon_loss = np.mean((output - X_batch) ** 2)
                epoch_loss += recon_loss
                batch_count += 1
                batch_elapsed = time.perf_counter() - batch_start_time
                logger.info(
                    "ART epoch %d/%d batch %d/%d: loss=%.4f, batch_time=%.2fs, samples/s=%.2f",
                    epoch + 1,
                    epochs,
                    batch_index + 1,
                    total_batches,
                    recon_loss,
                    batch_elapsed,
                    len(X_batch) / max(batch_elapsed, 1e-9),
                )
            
            avg_loss = epoch_loss / max(batch_count, 1)
            logger.info(
                "ART epoch %d/%d complete: avg_loss=%.4f, epoch_time=%.2fs",
                epoch + 1,
                epochs,
                avg_loss,
                time.perf_counter() - epoch_start,
            )
            
            # Store contemplation history
            self.contemplation_history.append(avg_loss)
    
    def transform(self, X: np.ndarray) -> np.ndarray:
        """Transform new data using learned references."""
        return self.forward(X)


# =============================================================================
# DISCRIMINATOR (Judges conv-likeness)
# =============================================================================

class Discriminator:
    """
    Discriminator network that judges if a transform is "conv-like".
    
    The discriminator compares reference-based transforms against
    ground-truth convolution outputs, providing adversarial feedback.
    """
    
    def __init__(self, input_channels: int, hidden_size: int = 64):
        self.W1 = np.random.randn(input_channels * 2, hidden_size) * 0.01
        self.b1 = np.zeros((1, hidden_size))
        self.W2 = np.random.randn(hidden_size, 1) * 0.01
        self.b2 = np.zeros((1, 1))
        
        # Binary cross-entropy state
        self.real_probability = 0.5
        self.fake_probability = 0.5
    
    def forward(self, real_features: np.ndarray, fake_features: np.ndarray) -> Tuple[float, float]:
        """
        Classify features as real (conv-like) or fake (reference-based).
        
        Returns:
            Tuple of (discriminator score for real, score for fake)
        """
        real_flat = real_features.reshape(real_features.shape[0], -1)
        fake_flat = fake_features.reshape(fake_features.shape[0], -1)

        # Combine real and fake for joint discrimination
        combined_real = np.concatenate([
            real_flat,
            fake_flat
        ], axis=1)
        
        # Simple scoring (in practice, would use full network)
        real_width = real_flat.shape[1]
        score_real = np.mean(combined_real[:, :real_width])
        score_fake = np.mean(combined_real[:, real_width:])
        
        return float(score_real), float(score_fake)
    
    def train_step(
        self,
        real: np.ndarray,
        fake: np.ndarray,
        learning_rate: float = 0.001
    ) -> Tuple[float, float]:
        """
        Single training step for the discriminator.
        
        Returns:
            Tuple of (loss, accuracy)
        """
        # Score real vs fake
        score_real, score_fake = self.forward(real, fake)
        
        # Compute simple accuracy
        accuracy = 0.5  # Placeholder for actual classification
        
        # Update weights (simplified)
        error = score_fake - score_real
        self.W1 += learning_rate * error * 0.01
        self.b1 += learning_rate * error * 0.001
        
        # Update probabilities
        self.real_probability = min(0.99, max(0.01, 0.5 + score_real * 0.1))
        self.fake_probability = min(0.99, max(0.01, 0.5 + score_fake * 0.1))
        
        loss = abs(error)
        accuracy = 1.0 / (1.0 + abs(error))
        
        return loss, accuracy


# =============================================================================
# MLP CLASSIFIER (Preserved from original)
# =============================================================================

class MLPClassifier:
    """MLP Classifier with SPDER activation (from previous update)."""
    
    def __init__(self, input_size: int, hidden_size: int, output_size: int, 
                 learning_rate: np.ndarray = None):
        self.W1 = np.random.randn(input_size, hidden_size) * 0.01
        self.b1 = np.zeros((1, hidden_size))
        self.W2 = np.random.randn(hidden_size, output_size) * 0.01
        self.b2 = np.zeros((1, output_size))
        self.lr = learning_rate if learning_rate is not None else np.array([0.01, 0.01, 0.01, 0.01])
    
    def spder(self, x: np.ndarray) -> np.ndarray:
        """SPDER: sin(x) * sqrt(|x|) - Periodic with damping."""
        return np.sin(x) * np.sqrt(np.abs(x) + 1e-8)
    
    def spder_derivative(self, x: np.ndarray) -> np.ndarray:
        """Derivative of SPDER."""
        abs_x = np.abs(x) + 1e-8
        sqrt_abs_x = np.sqrt(abs_x)
        sign_x = np.sign(x)
        sign_x = np.where(sign_x == 0, 1, sign_x)
        return sqrt_abs_x * np.cos(x) + (sign_x / (2 * sqrt_abs_x)) * np.sin(x)
    
    def softmax(self, x: np.ndarray) -> np.ndarray:
        exp_x = np.exp(x - np.max(x, axis=1, keepdims=True))
        return exp_x / np.sum(exp_x, axis=1, keepdims=True)
    
    def forward(self, X: np.ndarray) -> np.ndarray:
        X_norm = X / 127.5 - 1.0
        self.z1 = np.dot(X_norm, self.W1) + self.b1
        self.a1 = self.spder(self.z1)
        self.z2 = np.dot(self.a1, self.W2) + self.b2
        return self.softmax(self.z2)
    
    def backward(self, X: np.ndarray, y_true: np.ndarray, y_pred: np.ndarray):
        m = y_true.shape[0]
        X_norm = X / 127.5 - 1.0
        
        dz2 = y_pred - y_true
        dW2 = np.dot(self.a1.T, dz2) / m
        db2 = np.sum(dz2, axis=0, keepdims=True) / m
        
        da1 = np.dot(dz2, self.W2.T)
        dz1 = da1 * self.spder_derivative(self.z1)
        dW1 = np.dot(X_norm.T, dz1) / m
        db1 = np.sum(dz1, axis=0, keepdims=True) / m
        
        return dW1, db1, dW2, db2
    
    def update(self, X: np.ndarray, y_true: np.ndarray):
        y_pred = self.forward(X)
        dW1, db1, dW2, db2 = self.backward(X, y_true, y_pred)
        self.W1 -= self.lr[0] * dW1
        self.b1 -= self.lr[1] * db1
        self.W2 -= self.lr[2] * dW2
        self.b2 -= self.lr[3] * db2


# =============================================================================
# COMBINED PIPELINE: ART + MLP CLASSIFIER
# =============================================================================

class ART_MLP_Classifier:
    """
    Combined Adversarial Reference Transform + MLP Classifier.
    
    The ART layer preprocesses the input, then feeds to MLP for classification.
    """
    
    def __init__(
        self,
        input_shape: Tuple[int, int, int] = (28, 28, 1),
        reference_channels: int = 32,
        hidden_size: int = 100,
        output_size: int = 10
    ):
        # Adversarial Reference Transform
        self.art = AdversarialReferenceTransform(
            input_shape=input_shape,
            reference_channels=reference_channels,
            reference_grid_size=(3, 3),
            contemplation_depth=5
        )
        
        # Classification MLP
        self.mlp = MLPClassifier(
            input_size=np.prod(input_shape),
            hidden_size=hidden_size,
            output_size=output_size,
            learning_rate=np.array([0.01, 0.01, 0.01, 0.01])
        )
    
    def preprocess(self, X: np.ndarray) -> np.ndarray:
        """Apply ART preprocessing to input."""
        if len(X.shape) == 2:
            # Flatten to (B, H, W, C)
            size = int(np.sqrt(X.shape[1]))
            X = X.reshape(-1, size, size, 1)
        
        # Apply adversarial reference transform
        transformed = self.art.forward(X)
        
        # Flatten back
        return transformed.reshape(transformed.shape[0], -1)
    
    def fit(self, X_train: np.ndarray, y_train: np.ndarray, 
            art_epochs: int = 5, mlp_epochs: int = 100, batch_size: int = 100):
        """Train the combined pipeline."""
        logger.info("Phase 1: Training Adversarial Reference Transform...")
        self.art.fit(X_train, epochs=art_epochs, batch_size=batch_size)
        
        logger.info("Phase 2: Training MLP Classifier...")
        X_preprocessed = self.preprocess(X_train)
        
        for i in range(mlp_epochs):
            idx = np.random.randint(0, len(X_train), batch_size)
            X_batch = X_preprocessed[idx]
            y_batch = y_train[idx]
            self.mlp.update(X_batch, np.eye(10)[y_batch])
            
            if i % 1 == 0:
                acc = self.score(X_preprocessed, y_train)
                logger.info(f"MLP Epoch {i}: accuracy={acc:.4f}")
    
    def predict(self, X: np.ndarray) -> np.ndarray:
        X_proc = self.preprocess(X)
        return np.argmax(self.mlp.forward(X_proc), axis=1)
    
    def score(self, X: np.ndarray, y_true: np.ndarray) -> float:
        if X.ndim == 2:
            y_pred = np.argmax(self.mlp.forward(X), axis=1)
        else:
            y_pred = self.predict(X)
        return np.mean(y_pred == y_true)


# =============================================================================
# MAIN EXECUTION
# =============================================================================

if __name__ == "__main__":
    # Load data
    logger.info("Loading data...")
    X_train = read('../X_train.wav')[1].reshape(-1, 784)
    y_train = (read('../y_train.wav')[1] * 9).astype(int)
    X_test = read('../X_test.wav')[1].reshape(-1, 784)
    y_test = (read('../y_test.wav')[1] * 9).astype(int)
    
    # Reshape for ART
    X_train_img = X_train.reshape(-1, 28, 28, 1)
    X_test_img = X_test.reshape(-1, 28, 28, 1)
    
    # Initialize combined classifier
    logger.info("Initializing ART + MLP Classifier...")
    classifier = ART_MLP_Classifier(
        input_shape=(28, 28, 1),
        reference_channels=32,
        hidden_size=100,
        output_size=10
    )
    
    # Train
    classifier.fit(X_train_img, y_train, art_epochs=3, mlp_epochs=2, batch_size=100)
    
    # Evaluate
    accuracy = classifier.score(X_test_img, y_test)
    logger.info(f"Final Test Accuracy: {accuracy:.4f}")
    
    # Demo: Pure ART transform
    logger.info("Demonstrating standalone ART transform...")
    demo_input = X_test_img[:5]
    art_output = classifier.art.forward(demo_input)
    logger.info(f"Input shape: {demo_input.shape}")
    logger.info(f"ART output shape: {art_output.shape}")
    logger.info(f"Output range: [{art_output.min():.3f}, {art_output.max():.3f}]")