"""SOM class implementation"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['dist_fn', 'SOM', 'exp_sched', 'Scheduler']

# %% ../nbs/00_core.ipynb 2
import numpy as np
from sklearn.decomposition import PCA
from fastcore.all import *
import matplotlib.pyplot as plt

# %% ../nbs/00_core.ipynb 3
dist_fn = lambda x, ws: np.linalg.norm(ws-x, axis=-1)

# %% ../nbs/00_core.ipynb 4
class SOM:
    def __init__(self, 
                 grid_sz:tuple, # grid size
                 input_dim:int, # input dimension
                 init:str='random', # initialization method
                #  dist_fn=np.linalg.norm # distance function
                 dist_fn=dist_fn # distance function
                 )->None: # initialize the SOM
        "Initialize a Self-Organizing Map with given grid size and input dimension"
        store_attr()
        self.weights = None  # Will be initialized when fitting

# %% ../nbs/00_core.ipynb 5
@patch
def _initialize_weights_pca(self:SOM, X:np.ndarray)->np.ndarray:
    "Initialize weights using PCA of the input data"
    pca = PCA(n_components=2)
    pca.fit(X)
    
    # Create grid coordinates and scale by eigenvalues
    n = self.grid_sz[0]
    alpha = np.linspace(-1, 1, n) * np.sqrt(pca.explained_variance_[0])
    beta = np.linspace(-1, 1, n) * np.sqrt(pca.explained_variance_[1])
    
    # Create the grid
    alpha_grid, beta_grid = np.meshgrid(alpha, beta)
    
    # Initialize weights as linear combination of first two PCs
    return (alpha_grid[..., np.newaxis] * pca.components_[0] + 
            beta_grid[..., np.newaxis] * pca.components_[1])

# %% ../nbs/00_core.ipynb 6
@patch
def _initialize_weights(self:SOM, 
                        X:np.ndarray=None, # data matrix
                        method:str='random' # initialization method
                        )->np.ndarray: # weights matrix
    "Initialize weights using either random or PCA initialization"
    if method == 'random':
        return np.random.randn(*self.grid_sz, self.input_dim)
    elif method == 'pca':
        if X is None: 
            raise ValueError("Data matrix X required for PCA initialization")
        return self._initialize_weights_pca(X)

# %% ../nbs/00_core.ipynb 7
@patch
def _find_bmu(self:SOM, 
              x:np.ndarray, # input vector
              )->tuple: # coordinates of the best matching unit
    "Find coordinates of Best Matching Unit for input x"
    distances = self.dist_fn(x, self.weights)
    return np.unravel_index(np.argmin(distances), self.grid_sz)

# %% ../nbs/00_core.ipynb 8
@patch
def _grid_distances(self:SOM, 
                    bmu_pos:tuple # coordinates of the best matching unit
                    )->np.ndarray: # grid distances
    "Calculate grid distances from BMU position"
    rows_idx, cols_idx = np.ogrid[0:self.grid_sz[0], 0:self.grid_sz[1]]
    return (bmu_pos[0] - rows_idx)**2 + (bmu_pos[1] - cols_idx)**2

# %% ../nbs/00_core.ipynb 9
@patch
def _neighborhood_function(self:SOM, 
                           grid_dist:np.ndarray, # grid distances
                           sigma:float # neighborhood radius
                           )->np.ndarray: # neighborhood function values
    "Calculate neighborhood function values"
    return np.exp(-grid_dist/(2*sigma**2))

# %% ../nbs/00_core.ipynb 10
@patch
def _update_weights(self:SOM, 
                    x:np.ndarray, # input vector
                    learning_rate:float, # learning rate
                    sigma:float # neighborhood radius
                    )->None: # update weights
    "Update weights for a single input vector"
    bmu_pos = self._find_bmu(x)
    grid_dist = self._grid_distances(bmu_pos)
    neighborhood = self._neighborhood_function(grid_dist, sigma)
    self.weights += learning_rate * neighborhood[..., np.newaxis] * (x - self.weights)

# %% ../nbs/00_core.ipynb 11
def exp_sched(start_val, end_val, i, n_steps):
    decay = -np.log(end_val/start_val)/n_steps
    return start_val * np.exp(-decay * i)

# %% ../nbs/00_core.ipynb 12
class Scheduler:
    def __init__(self, start_val, end_val, step_size, 
                 n_samples, n_epochs, decay_fn=exp_sched):
        store_attr()
        self.current_step = 0
        self.current_value = start_val
        self.total_steps = (n_samples * n_epochs) // step_size
    
    def step(self, total_samples):
        if total_samples % self.step_size == 0:
            self.current_value = self.decay_fn(
                self.start_val, self.end_val, 
                self.current_step, self.total_steps
            )
            self.current_step += 1
        return self.current_value

# %% ../nbs/00_core.ipynb 13
@patch
def fit(self:SOM, 
        X:np.ndarray, # data matrix
        n_epochs:int=20, # number of epochs
        lr_scheduler:Scheduler=None, # learning rate scheduler
        sigma_scheduler:Scheduler=None, # neighborhood radius scheduler
        shuffle:bool=True, # shuffle data
        verbose:bool=True # verbose
        )->tuple: # tuple of weights, quantization errors, and topographic errors
    "Train the SOM on input data X"
    # Initialize weights if not already done
    if self.weights is None:
        self.weights = self._initialize_weights(X, self.init)
    
    # Setup default schedulers if none provided
    if lr_scheduler is None:
        lr_scheduler = Scheduler(1.0, 0.01, 100, len(X), n_epochs)
    if sigma_scheduler is None:
        sigma_scheduler = Scheduler(max(self.grid_sz)/2, 1.0, 100, len(X), n_epochs)
    
    qe_errors, te_errors = [], []
    for epoch in range(n_epochs):
        X_ = np.random.permutation(X) if shuffle else X.copy()
        
        # Train on each input vector
        for i, x in enumerate(X_):
            total_samples = epoch * len(X) + i
            lr = lr_scheduler.step(total_samples)
            sigma = sigma_scheduler.step(total_samples)
            self._update_weights(x, lr, sigma)
        
        # Calculate errors
        qe = self.quantization_error(X)
        te = self.topographic_error(X)
        qe_errors.append(qe)
        te_errors.append(te)
        
        if verbose:
            print(f'Epoch: {epoch+1} | QE: {qe:.4f}, TE: {te:.4f}')
    
    return self.weights, qe_errors, te_errors

# %% ../nbs/00_core.ipynb 14
@patch
def quantization_error(self:SOM, 
                       X:np.ndarray # data matrix
                       )->float: # quantization error
    "Calculate average distance between each input vector and its BMU"
    return np.array([
        self.dist_fn(x, self.weights).min() 
        for x in X
    ]).mean()

# %% ../nbs/00_core.ipynb 15
@patch
def topographic_error(self:SOM, 
                      X:np.ndarray # data matrix
                      )->float: # topographic error
    "Calculate proportion of data vectors where 1st and 2nd BMUs are not adjacent"
    def _check_bmu_adjacency(x):
        # Get indices of two best matching units
        distances = self.dist_fn(x, self.weights)
        flat_indices = np.argpartition(distances.flatten(), 2)[:2]
        indices = np.unravel_index(flat_indices, self.grid_sz)
        # Check if any coordinate differs by more than 1
        return any(np.abs(x-y) > 1 for x,y in indices)
    
    n_errors = sum(_check_bmu_adjacency(x) for x in X)
    return 100 * n_errors / len(X)  # Return percentage

# %% ../nbs/00_core.ipynb 16
@patch
def transform(self:SOM, 
              X:np.ndarray # data matrix
              )->np.ndarray: # bmu coordinates
    "Find Best Matching Unit (BMU) coordinates for each input vector"
    bmu_coords = np.zeros((len(X), 2), dtype=int)
    for i, x in enumerate(X):
        bmu_coords[i] = self._find_bmu(x)
    return bmu_coords

# %% ../nbs/00_core.ipynb 17
@patch  
def predict(self:SOM, 
            X:np.ndarray # data matrix
            )->np.ndarray: # bmu coordinates
    "Alias for transform method to follow sklearn convention"
    return self.transform(X)

# %% ../nbs/00_core.ipynb 18
@patch  
def _calculate_umatrix(self:SOM)->np.ndarray:
    "Calculate U-Matrix values for current weights"
    def _neighbor_distances(pos):
        # Offsets for 8 neighbors
        nbr_offsets = [
            (-1,-1), (-1,0), (-1,1),  # top-left, top, top-right
            (0,-1),          (0,1),   # left, right
            (1,-1),  (1,0),  (1,1)    # bottom-left, bottom, bottom-right
        ]
    
        distances = []
        weights = []
        for dr, dc in nbr_offsets:
            r, c = pos
            nbr_r, nbr_c = r+dr, c+dc
            if (nbr_r >= 0 and nbr_r < self.grid_sz[0] and 
                nbr_c >= 0 and nbr_c < self.grid_sz[1]):
                w = 1/np.sqrt(dr**2 + dc**2)  # weight by distance
                weights.append(w)
                d = self.dist_fn(self.weights[r,c], self.weights[nbr_r,nbr_c])
                distances.append(d)
        return np.average(distances, weights=weights)
        
    umatrix = np.zeros(self.grid_sz)
    for i, j in np.ndindex(self.grid_sz):
        umatrix[i,j] = _neighbor_distances((i,j))
    return umatrix

# %% ../nbs/00_core.ipynb 19
@patch
def umatrix(self:SOM)->np.ndarray:
    "Return the U-Matrix for current weights"
    return self._calculate_umatrix()

# %% ../nbs/00_core.ipynb 20
@patch
def plot_umatrix(self:SOM, 
                 figsize=(8,6), 
                 cmap='viridis_r'):
    "Plot U-Matrix visualization"
    plt.figure(figsize=figsize)
    plt.imshow(self.umatrix(), cmap=cmap, interpolation='nearest')
    plt.colorbar(label='Average distance to neighbors')
    plt.title('U-Matrix')
    plt.tight_layout();
