import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.distance import cdist
import scipy.io

class NeuralGasNetwork:
    def __init__(self, X, params, plot_flag=False):
        self.X = X
        self.params = params
        self.plot_flag = plot_flag

        self.n_data, self.n_dim = X.shape
        self.N = params["N"]
        self.MaxIt = params["MaxIt"]
        self.tmax = params["tmax"]
        self.epsilon_initial = params["epsilon_initial"]
        self.epsilon_final = params["epsilon_final"]
        self.lambda_initial = params["lambda_initial"]
        self.lambda_final = params["lambda_final"]
        self.T_initial = params["T_initial"]
        self.T_final = params["T_final"]

        self.initialize_network()

    def initialize_network(self):
        np.random.shuffle(self.X)  # Randomize the input data order
        Xmin, Xmax = self.X.min(axis=0), self.X.max(axis=0)
        self.w = np.random.uniform(Xmin, Xmax, (self.N, self.n_dim))  # Initialize weights randomly
        self.C = np.zeros((self.N, self.N), dtype=int)  # Connectivity matrix
        self.t = np.zeros((self.N, self.N), dtype=int)  # Age matrix
        self.tt = 0  # Time step counter

    def train(self):
        for it in range(self.MaxIt):
            for l in range(self.n_data):
                x = self.X[l, :]

                # Compute distances & sort units
                d = cdist([x], self.w, metric='euclidean').flatten()
                sort_order = np.argsort(d)

                # Compute learning parameters
                epsilon = self.epsilon_initial * (self.epsilon_final / self.epsilon_initial) ** (self.tt / self.tmax)
                lambda_ = self.lambda_initial * (self.lambda_final / self.lambda_initial) ** (self.tt / self.tmax)
                T = self.T_initial * (self.T_final / self.T_initial) ** (self.tt / self.tmax)

                # Adaptation step
                for ki in range(self.N):
                    i = sort_order[ki]
                    self.w[i, :] += epsilon * np.exp(-ki / lambda_) * (x - self.w[i, :])

                self.tt += 1  # Increment global time step

                # Creating links
                i, j = sort_order[:2]  # Best matching units
                self.C[i, j] = self.C[j, i] = 1
                self.t[i, j] = self.t[j, i] = 0

                # Aging links
                self.t[i, :] += 1
                self.t[:, i] += 1

                # Remove old links
                old_links = self.t[i, :] > T
                self.C[i, old_links] = 0
                self.C[old_links, i] = 0

            # Plot results if enabled
            if self.plot_flag:
                self.plot_results()

        return {"w": self.w, "C": self.C, "t": self.t}

    def plot_results(self):
        plt.figure(figsize=(6, 6))
        plt.scatter(self.X[:, 0], self.X[:, 1], c='gray', alpha=0.5, label="Data")
        plt.scatter(self.w[:, 0], self.w[:, 1], c='red', label="Neurons")

        for i in range(self.N):
            for j in range(i + 1, self.N):
                if self.C[i, j] == 1:
                    plt.plot([self.w[i, 0], self.w[j, 0]], [self.w[i, 1], self.w[j, 1]], 'b-')

        plt.legend()
        plt.title("Neural Gas Network")
        plt.show()

# Example usage
if __name__ == "__main__":
    
    data = scipy.io.loadmat('HWMar26data.mat')
    X=data.get('A')

    params = {
        "N": 20,
        "MaxIt": 30,
        "tmax": 9000,
        "epsilon_initial": 0.2,
        "epsilon_final": 0.001,
        "lambda_initial": 0.75,
        "lambda_final": 0.1,
        "T_initial": 5,
        "T_final": 30
    }

    ngn = NeuralGasNetwork(X, params, plot_flag=True)
    result = ngn.train()